Feature expansion and bug-fix

.gitignore

@@ -106,4 +106,4 @@ ENV/
 *.sublime-project
 *.sublime-workspace
 
-generated/
+generated/

dwave/system/testing.py

@@ -28,9 +28,15 @@
 class MockDWaveSampler(dimod.Sampler, dimod.Structured):
     """Mock sampler modeled after DWaveSampler that can be used for tests.
 
-    Properties and topology parameters are populated qualitatively matching
+    Properties and topology parameters are populated to qualitatively match
     online systems, and a placeholder sampler routine based on steepest descent
-    is instantiated.
+    is instantiated by default.
+
+    The `EXACT_SOLVER_CUTOFF_DEFAULT` defines the problem size threshold for using the exact solver.
+    For problems with fewer variables than this threshold, the exact ground state is computed 
+    using a brute-force solver. This provides a reproducible solution for small problem sizes.
+
+    For larger problems, the `SteepestDescentSampler` is used as a placeholder solver.
 
     Args:
         nodelist (iterable of ints, optional):
@@ -73,6 +79,20 @@ class MockDWaveSampler(dimod.Sampler, dimod.Structured):
             parameters. By default ``initial_state`` can also be mocked, if
             dwave-greedy is installed. All other parameters are ignored and a 
             warning will be raised by default.
+
+        substitute_sampler (dimod.Sampler, optional, default=SteepestDescentSampler()):
+            The sampler to be used as a substitute when executing the mock sampler. 
+            By default, `SteepestDescentSampler()` is employed, which performs a 
+            deterministic steepest descent optimization on the BQM. Supported options are
+            any dimod-compatible sampler to customize the sampling behavior of 
+            `MockDWaveSampler()`.
+
+        substitute_kwargs (dict, optional, default=[]):
+            A dictionary of keyword arguments to pass to the `substitute_sampler`'s 
+            `sample` method. This allows users to configure the substitute sampler 
+            with specific parameters like `num_reads`, `initial_state`, or other 
+            sampler-specific options. If not provided, an empty dictionary is used 
+            by default.
 
         exact_solver_cutoff (int, optional, default=:attr:`EXACT_SOLVER_CUTOFF_DEFAULT`):
             For problems smaller or equal in size to ``exact_solver_cutoff``, the
@@ -108,23 +128,38 @@ class MockDWaveSampler(dimod.Sampler, dimod.Structured):
         -1
 
     """
-    # Feature suggestion - add seed as an optional input, to allow reproducibility.
-
     nodelist = None
     edgelist = None
     properties = None
     parameters = None
-
-    # by default, use ExactSolver for problems up to size (inclusive):
+
     EXACT_SOLVER_CUTOFF_DEFAULT = 16
 
     def __init__(self,
                  nodelist=None, edgelist=None, properties=None,
                  broken_nodes=None, broken_edges=None,
                  topology_type=None, topology_shape=None,
                  parameter_warnings=True,
+                 substitute_sampler=None, 
+                 substitute_kwargs=None, 
                  exact_solver_cutoff=EXACT_SOLVER_CUTOFF_DEFAULT,
                  **config):
+
+        self.mocked_parameters={'answer_mode',
+                       'max_answers',
+                       'num_reads',
+                       'label',
+                       'initial_state',
+        }
+
+        if substitute_sampler is None:
+            substitute_sampler = SteepestDescentSampler()
+        self.substitute_sampler = substitute_sampler
+
+        if substitute_kwargs is None:
+            substitute_kwargs = {} 
+        self.substitute_kwargs = substitute_kwargs
+
         self.parameter_warnings = parameter_warnings
         self.exact_solver_cutoff = exact_solver_cutoff
 
@@ -158,7 +193,6 @@ def __init__(self,
             'chip_id': 'MockDWaveSampler',
             'topology': {'type': topology_type, 'shape': topology_shape}
         }
-
         #Create graph object, introduce defects per input arguments
         if nodelist is not None:
             self.nodelist = nodelist.copy()
@@ -171,6 +205,11 @@ def __init__(self,
                                  self.properties['topology']['shape'],
                                  self.nodelist, self.edgelist)
 
+        if topology_type == 'pegasus':
+            m = self.properties['topology']['shape'][0]
+            num_qubits = m*(m-1)*24  # fabric_only=True technicality
+        else:
+            num_qubits = len(solver_graph)
         if broken_nodes is None and broken_edges is None:
             self.nodelist = sorted(solver_graph.nodes)
             self.edgelist = sorted(tuple(sorted(edge))
@@ -189,7 +228,7 @@ def __init__(self,
                                    and (v, u) not in broken_edges)
         #Finalize yield-dependent properties:
         self.properties.update({
-            'num_qubits': len(solver_graph),
+            'num_qubits': num_qubits,
             'qubits': self.nodelist.copy(),
             'couplers': self.edgelist.copy(),
             'anneal_offset_ranges': [[-0.5, 0.5] if i in self.nodelist
@@ -295,6 +334,7 @@ def __init__(self,
             'tags': [],
             'category': 'qpu',
             'quota_conversion_rate': 1,
+            'fast_anneal_time_range': [0.005, 83000.0],
         })
 
         if properties is not None:
@@ -313,15 +353,9 @@ def from_qpu_sampler(cls, sampler):
     def sample(self, bqm, **kwargs):
 
         # Check kwargs compatibility with parameters and substitute sampler:
-        mocked_parameters={'answer_mode',
-                           'max_answers',
-                           'num_reads',
-                           'label',
-                           'initial_state',
-                           }
         for kw in kwargs:
             if kw in self.parameters:
-                if self.parameter_warnings and kw not in mocked_parameters:
+                if self.parameter_warnings and kw not in self.mocked_parameters:
                     warnings.warn(f'{kw!r} parameter is valid for DWaveSampler(), '
                                   'but not mocked in MockDWaveSampler().')
             else:
@@ -344,19 +378,21 @@ def sample(self, bqm, **kwargs):
         label = kwargs.get('label')
         if label is not None:
             info.update(problem_label=label)
-
-        #Special handling of flux_biases, for compatibility with virtual graphs
-
+
+        # Special handling of flux_biases, for compatibility with virtual graphs
         flux_biases = kwargs.get('flux_biases')
         if flux_biases is not None:
             self.flux_biases_flag = True
 
-        substitute_kwargs = {'num_reads' : kwargs.get('num_reads')}
-        if substitute_kwargs['num_reads'] is None:
-            substitute_kwargs['num_reads'] = 1
+        # Create a local dictionary combining self.substitute_kwargs and relevant kwargs
+        substitute_kwargs = self.substitute_kwargs.copy()
+
+        # Handle 'num_reads', defaulting to 1 if not provided
+        num_reads = kwargs.get('num_reads', substitute_kwargs.get('num_reads', 1))
+        substitute_kwargs['num_reads'] = num_reads
 
-        initial_state = kwargs.get('initial_state')
-        if initial_state is not None:
+        if 'initial_state' in kwargs:
+            initial_state = kwargs['initial_state']
             # Initial state format is a list of (qubit,values)
             # value=3 denotes an unused variable (should be absent
             # from bqm). 
@@ -366,9 +402,12 @@ def sample(self, bqm, **kwargs):
                           if pair[1]!=3],dtype=float),
                 [pair[0] for pair in initial_state if pair[1]!=3])
 
-        ss = SteepestDescentSampler().sample(bqm, **substitute_kwargs)
-        ss.info.update(info)
+        sampler_kwargs = kwargs.copy()
+        sampler_kwargs.update(substitute_kwargs)
 
+        ss = self.substitute_sampler.sample(bqm, **sampler_kwargs)
+        ss.info.update(info)
+
         # determine ground state exactly for small problems
         if 0 < len(bqm) <= self.exact_solver_cutoff and len(ss) >= 1:
             ground = dimod.ExactSolver().sample(bqm).truncate(1)

tests/test_mock_sampler.py

@@ -216,7 +216,86 @@ def test_yield_arguments(self):
                                    broken_edges=delete_edges)
         self.assertTrue(len(sampler.nodelist)==4)
         self.assertTrue(len(sampler.edgelist)==1)
+
+    def test_custom_mock_sampler(self):
+        """Test that MockDWaveSampler uses the provided custom mocking_sampler."""
+
+        # Define a constant sampler that always returns the same sample
+        class ConstantSampler(dimod.Sampler):
+            properties = {}
+            parameters = {'num_reads': []}
+
+            def sample(self, bqm, **kwargs):
+                num_reads = kwargs.get('num_reads', 1)
+                sample = {v: 1 for v in bqm.variables}
+                samples = [sample] * num_reads
+                energies = [bqm.energy(sample) for sample in samples]
+                return dimod.SampleSet.from_samples(samples, vartype=bqm.vartype, energy=energies)
+
+        constant_sampler = ConstantSampler()
+
+        # Create a simple BQM
+        bqm = dimod.BQM({'a': 1, 'b': 1}, {}, 0.0, vartype="SPIN")
+
+        # Instantiate MockDWaveSampler with nodelist and edgelist including 'a' and 'b'
+        sampler = MockDWaveSampler(
+            substitute_sampler=constant_sampler,
+            nodelist=['a', 'b'],
+            edgelist=[('a', 'b')]
+        )
+
+        # Sample using the MockDWaveSampler with the custom sampler
+        ss = sampler.sample(bqm, num_reads=2, answer_mode='raw')  # First sample is overwritten by ExactSampler
 
+        # Reconstruct the second sample as a dictionary
+        second_sample_array = ss.record.sample[1]
+        variables = ss.variables
+        second_sample = dict(zip(variables, second_sample_array))
+
+        # Expected sample from the custom sampler
+        expected_sample = {v: 1 for v in bqm.variables}
+
+        # Check that the sample returned is as expected from the custom sampler
+        self.assertEqual(second_sample, expected_sample, 'Second sample was not the expected excited state')
+
+    def test_mocking_sampler_params(self):
+        """Test that mocking_sampler_params are correctly passed to the mocking_sampler."""
+
+        # Define a constant sampler that checks for a custom parameter
+        class ConstantSampler(dimod.Sampler):
+            properties = {}
+            parameters = {'custom_param': []}
+
+            def sample(self, bqm, **kwargs):
+                custom_param = kwargs.get('custom_param')
+                # Raise exception if parameters passed incorrectly
+                if custom_param != 'test_value':
+                    raise ValueError("custom_param not passed correctly")
+                # Return a default sample
+                sample = {v: -1 for v in bqm.variables}
+                return dimod.SampleSet.from_samples_bqm(sample, bqm)
+
+        constant_sampler = ConstantSampler()
+
+        # Create a simple BQM
+        bqm = dimod.BQM({'a': 1, 'b': 1}, {('a', 'b'): 1}, 0.0, vartype="SPIN")
+
+        # Instantiate MockDWaveSampler with nodelist and edgelist including 'a' and 'b'
+        sampler = MockDWaveSampler(
+            substitute_sampler=constant_sampler,
+            substitute_kwargs={'custom_param': 'test_value'},
+            nodelist=['a', 'b'],
+            edgelist=[('a', 'b')]
+        )
+
+        # Sample using the MockDWaveSampler
+        ss = sampler.sample(bqm)
+
+        # Check that the sample returned is as expected from the custom sampler
+        expected_sample = {'a': -1, 'b': -1}
+        self.assertEqual(ss.first.sample, expected_sample)
+        self.assertEqual(ss.first.energy, bqm.energy(expected_sample))
+
 class TestMockLeapHybridDQMSampler(unittest.TestCase):
     def test_sampler(self):
         sampler = MockLeapHybridDQMSampler()

tests/test_virtual_graph_composite.py

@@ -33,7 +33,7 @@ def test_smoke(self):
 
         # depending on how recenlty flux bias data was gathered, this may be true
         child_sampler.flux_biases_flag = False
-
+        child_sampler.mocked_parameters.add('flux_biases')  # Don't raise warning
         if sampler.flux_biases:
             sampler.sample_ising({'a': -1}, {})
             self.assertTrue(child_sampler.flux_biases_flag)  # true when some have been provided to sample_ising