Update Adding-constrained-MH with Feroz' recent changes

docs/source/example.md

@@ -1,7 +1,7 @@
 # Example 2D Ising Model
 
 
-##### 1. Initialise an energy model
+<!-- ##### 1. Initialise an energy model
 In this step, we define a classical energy function over binary spin configurations.
 This energy model is the target distribution that the QeMCMC sampler will explore.
 
@@ -45,7 +45,7 @@ At each MCMC step, a subgroup is sampled according to subgroup_probs.
 ##### 3. Create and run QeMCMC
 Finally, we initialise the quantum-enhanced Markov chain and generate a single proposal using simulated quantum time evolution.
 ```python
-from qemcmc.qemcmc import QeMCMC
+from qemcmc import QeMCMC
 
 sampler = QeMCMC(
     model=model,
@@ -60,4 +60,4 @@ s = "01010"
 s_prime = sampler.get_s_prime(s)
 print("proposal:", s, "->", s_prime)
 ```
-Here, each call to get_s_prime runs a quantum circuit to generate a proposal state conditioned on the current configuration.
+Here, each call to get_s_prime runs a quantum circuit to generate a proposal state conditioned on the current configuration. -->

src/qemcmc/circuits.py

@@ -3,6 +3,7 @@
 from qemcmc.model.energy_model import EnergyModel
 from typing import List
 
+
 class CircuitMaker:
     """
     Constructs and simulates quantum circuits used to generate QeMCMC proposals.
@@ -38,14 +39,11 @@ def __init__(self, model: EnergyModel, delta_time: float = 0.8):
         self.model = model
         self.delta_time = delta_time
         self.n_qubits = model.n
-
-
         self.dev = qml.device("lightning.qubit", wires=self.n_qubits)
         #self.dev = qml.device("default.tensor", wires=self.n_qubits, method="mps", max_bond_dim=2, contract="auto-mps")
 
         self.model_type = model.model_type
-        # cache devices for dynamic subgroup sizes if needed
-        self.devices = {}
+        self.devices = {}  # cache devices for dynamic subgroup sizes if needed
 
     def _get_device(self, num_wires: int):
         """Get or create a PennyLane device for the given number of wires."""
@@ -98,7 +96,6 @@ def get_problem_hamiltonian(self, couplings: List[np.ndarray], sign: int = 1) ->
                     for q in index_tuple[1:]:
                         term = term @ qml.PauliZ(q)
                     total_hamiltonian += (sign * spin_sign * coeff) * term
-
 
                 elif self.model_type == "binary":
                     # 0.5 * (I - Z) for first variable
@@ -117,9 +114,9 @@ def get_problem_hamiltonian(self, couplings: List[np.ndarray], sign: int = 1) ->
     def get_mixer_hamiltonian(self, num_wires: int = None) -> qml.Hamiltonian:
         """
         Constructs the Mixer Hamiltonian: Σ X_i.
-        
+
         This can be for the full system or a subgroup.
-        
+
         Parameters
         ----------
         num_wires : int, optional
@@ -158,45 +155,41 @@ def get_state_vector(self, s: str, weights: List[float], time: float, mix_weight
         Notes
         -----
         The total Hamiltonian simulated by the circuit is a weighted sum of the problem Hamiltonian terms and the mixer Hamiltonian:
-        
+
         In Ferguson et al. (2025) [arXiv:2506.19538], we use gammas to weight the entire problem Hamiltonian vs the mixer vs the constraint Hamiltonian, such that the total Hamiltonian is:
-        
+
         H = g_p * H_p + g_m * H_m + g_c * H_c
 
         but here we allow for separate weights for each coupling tensor term, as well as a separate gamma for the mixer. The total Hamiltonian is then:
 
         H = (w_b1*H_b1+ w_b2*H_b2 +...+w_b2*H_bm) + g_m * H_m
 
-        In other words, the constraint hamiltonian is absorbed in the coupling list, and weighted by the corresponding gamma in the gammas list. 
+        In other words, the constraint hamiltonian is absorbed in the coupling list, and weighted by the corresponding gamma in the gammas list.
         This allows for more flexible weighting of different terms, and also allows us to use the same code for both constrained and unconstrained problems (by simply including or excluding the constraint Hamiltonian in the coupling list and adjusting the gammas accordingly).
-        
+
         Note that it is assumed that each term is already normalised appropriately, so the gammas can be interpreted as the relative weights of each term in the total Hamiltonian.
 
         """
-        
+
         num_wires = len(s)
         dev = self._get_device(num_wires)
-        
+
         if mix_weight < 0 or mix_weight > 1:
             raise ValueError(f"mix_weight must be between 0 and 1. Got {mix_weight}")
-
-
+
         if np.any(np.array(weights) < 0):
             raise ValueError(f"Weights must be non-negative. Got {weights}")
-        
+
         self.time = time
         self.num_trotter_steps = int(np.floor((self.time / self.delta_time)))
-
-        coeff_mixer = mix_weight
-        coeff_problem = weights 
-
 
+        coeff_mixer = mix_weight
+        coeff_problem = weights
 
         H_total = qml.Hamiltonian(
-            [coeff_mixer] + list(np.ones(len(self.model.normalised_couplings))), 
-            [self.get_mixer_hamiltonian(num_wires)] + 
-            [self.get_problem_hamiltonian(couplings=[self.model.normalised_couplings[i]], sign=coeff_problem[i]) 
-             for i in range(len(self.model.normalised_couplings))]
+            [coeff_mixer] + list(np.ones(len(self.model.normalised_couplings))),
+            [self.get_mixer_hamiltonian(num_wires)]
+            + [self.get_problem_hamiltonian(couplings=[self.model.normalised_couplings[i]], sign=coeff_problem[i]) for i in range(len(self.model.normalised_couplings))],
         )
 
         @qml.qnode(dev)
@@ -237,30 +230,30 @@ def get_sample(self, s_cg: str, time: float, mix_weight: float, local_couplings:
 
         num_wires = len(s_cg)
         dev = self._get_device(num_wires)
-        
+
         if mix_weight < 0 or mix_weight > 1:
             raise ValueError(f"mix_weight must be between 0 and 1. Got {mix_weight}")
-        
+
         if np.any(np.array(weights) < 0):
             raise ValueError(f"Weights must be non-negative. Got {weights}")
-        
+
         num_trotter_steps = int(np.floor((time / self.delta_time)))
-        
+
         coeff_mixer = mix_weight
         coeff_problem = weights
 
-        #coeff_problem = [coeff for i, coeff in enumerate(coeff_problem) if local_couplings[i].ndim > 0 and np.any(local_couplings[i] != 0)]
-        #local_couplings = [coupling for coupling in local_couplings if coupling.ndim > 0 and np.any(coupling != 0)]
+        # coeff_problem = [coeff for i, coeff in enumerate(coeff_problem) if local_couplings[i].ndim > 0 and np.any(local_couplings[i] != 0)]
+        # local_couplings = [coupling for coupling in local_couplings if coupling.ndim > 0 and np.any(coupling != 0)]
         if len(local_couplings) == 0:
             # If there are no local couplings, just return the input state
             print("no non-zero local couplings, skipping evolution and returning input state")
             return s_cg
-        
+
         H_total = qml.Hamiltonian(
-            [coeff_mixer] + list(np.ones(len(local_couplings))), 
-            [self.get_mixer_hamiltonian(num_wires)] + 
-            [self.get_problem_hamiltonian(couplings=[local_couplings[i]], sign=coeff_problem[i]) for i in range(len(local_couplings))]
+            [coeff_mixer] + list(np.ones(len(local_couplings))),
+            [self.get_mixer_hamiltonian(num_wires)] + [self.get_problem_hamiltonian(couplings=[local_couplings[i]], sign=coeff_problem[i]) for i in range(len(local_couplings))],
         )
+
         # set qnode to use our device with dynamically chosen wires
         @qml.qnode(dev, shots=1)
         def quantum_evolution(input_string):
@@ -317,9 +310,9 @@ def update(self, s: str, subgroup_choice: List[int], local_couplings: list, gamm
         str
             The updated bitstring s'.
         """
-        
+
         self._assert_bitstring(s)
-        
+
         # Get s_cg' for the subgroup and reconstruct full s' using s and s_cg'
         s_cg = "".join([s[i] for i in subgroup_choice])
         s_cg_prime = self.get_sample(s_cg, time, gamma, local_couplings)

src/qemcmc/coarse_grain.py

@@ -1,30 +1,26 @@
-# Internal
 from qemcmc.utils.helpers import validate_subgroups
-
-# External
 import numpy as np
 
 
 class CoarseGraining:
-    def __init__(self, n, subgroups=None, subgroup_probs=None, repeated= True):
-        """
-        CoarseGraining class to generate partitions of spins for quantum proposals.
-
-        Parameters
-        ----------
-        n : int
-            Number of spins in the system.
-        subgroups : list[list[int]], optional
-            A list of subgroups, where each subgroup is a list of spin indices. If None, the entire set of spins is treated as one subgroup.
-        subgroup_probs : list[float], optional
-            A list of probabilities corresponding to each subgroup, used for weighted random selection. Must sum to 1. If None, subgroups are selected uniformly at random.
-        repeated : bool, optional
-            If True, then multiple subgroups are run on the quantum computer in serial, if not then only one subgroup is selected at random and run on the quantum computer. Default is True.
+    """
+    CoarseGraining class to generate partitions of spins for quantum proposals.
 
+    Parameters
+    ----------
+    n : int
+        Number of spins in the system.
+    subgroups : list[list[int]], optional
+        A list of subgroups, where each subgroup is a list of spin indices. If None, the entire set of spins is treated as one subgroup.
+    subgroup_probs : list[float], optional
+        A list of probabilities corresponding to each subgroup, used for weighted random selection. Must sum to 1. If None, subgroups are selected uniformly at random.
+    repeated : bool, optional
+        If True, then multiple subgroups are run on the quantum computer in serial, if not then only one subgroup is selected at random and run on the quantum computer. Default is True.
+    """
 
-        """
-
+    def __init__(self, n, subgroups=None, subgroup_probs=None, repeated=True):
 
+        self._user_specified = subgroups is not None
 
         if subgroups is None:
             subgroups = [list(range(n))]
@@ -40,22 +36,34 @@ def __init__(self, n, subgroups=None, subgroup_probs=None, repeated= True):
         self.subgroup_probs = subgroup_probs
         self.repeated = repeated
 
-
-    def sample(self, rng=np.random):
-        idx = rng.choice(len(self.subgroups), p=self.subgroup_probs)
+    def sample(self) -> list[int]:
+        """
+        Randomly samples a subgroup according to the specified probability distribution.
+        """
+        idx = np.random.choice(len(self.subgroups), p=self.subgroup_probs)
         return self.subgroups[idx]
 
-    def get_partitions(self, m: int, rng=np.random):
+    def get_partitions(self, m: int) -> list[list[int]]:
         """
-        Randomly partitions the n spins into m disjoint subgroups.
-        Sizes will be approximately n/m.
+        Returns partitions of spins for sequential quantum updates.
+
+        If the user provided explicit subgroups at initialisation, those are
+        returned directly (all if ``repeated=True``, otherwise just the first).
+        If no subgroups were specified, random disjoint partitions of
+        approximate size n/m are generated.
         """
+
+        if self._user_specified:
+            if self.repeated:
+                return self.subgroups
+            else:
+                return [self.subgroups[0]]
+
         spins = np.arange(self.n)
-        rng.shuffle(spins)
+        np.random.shuffle(spins)
         # array_split handles uneven divisions gracefully
         chunks = [list(chunk) for chunk in np.array_split(spins, m)]
         if self.repeated:
             return chunks
         else:
-            return [chunks[0],]
-
+            return [chunks[0]]