From 5b85102bf5c87af27753899e6b5e4fd2f269de0b Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:41:27 +0100
Subject: [PATCH 01/12] docs: fix qemcmc.qemcmc module import issue

---
 docs/source/example.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/source/example.md b/docs/source/example.md
index 0d65a5c..1953d2b 100644
--- a/docs/source/example.md
+++ b/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. -->

From 970d904f278ecd2a1957cb589a058c9fc784a61b Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:55:52 +0100
Subject: [PATCH 02/12] refactor: fix imports

---
 src/qemcmc/sampler/classical_proposal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qemcmc/sampler/classical_proposal.py b/src/qemcmc/sampler/classical_proposal.py
index 8352d20..b10589f 100644
--- a/src/qemcmc/sampler/classical_proposal.py
+++ b/src/qemcmc/sampler/classical_proposal.py
@@ -10,7 +10,7 @@ class ClassicalProposal(Proposal):
 
     This class implements purely classical proposal mechanisms for MCMC.
     New candidate states are generated either by sampling a completely
-    random (uniform) configuration, or by performing a local single-spin or 
+    random (uniform) configuration, or by performing a local single-spin or
     two-spin flip.
 
     Parameters

From 717e5947e13a8ebc454c25e4985b487bc054ea73 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:56:10 +0100
Subject: [PATCH 03/12] fix: make class abstract

---
 src/qemcmc/sampler/proposal.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/qemcmc/sampler/proposal.py b/src/qemcmc/sampler/proposal.py
index 4c9cb3d..ac5aab3 100644
--- a/src/qemcmc/sampler/proposal.py
+++ b/src/qemcmc/sampler/proposal.py
@@ -1,13 +1,11 @@
-from abc import abstractmethod
-from typing import Optional
-
+from abc import abstractmethod, ABC
 from qemcmc.model import EnergyModel
 from qemcmc.utils import MCMCState
 
 
-class Proposal:
+class Proposal(ABC):
     """
-    Base class for producing proposals for Markov Chain Monte Carlo algorithms.
+    Abstract base class for producing proposals for Markov Chain Monte Carlo algorithms.
 
     Subclasses implement the proposal mechanism by defining an
     ``update(state)`` method that generates a candidate state from the current one
@@ -30,7 +28,7 @@ def update(self, state: MCMCState) -> MCMCState:
 
         This method should be implemented by subclasses to define the specific proposal strategy
         (e.g., single-spin flips, block updates, or quantum proposals).
-        
+
         Parameters
         ----------
         state : MCMCState
@@ -41,4 +39,4 @@ def update(self, state: MCMCState) -> MCMCState:
         MCMCState
             A new candidate state.
         """
-        pass
\ No newline at end of file
+        pass

From a0416829a69e898cf2e10be7af799de123618338 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:57:43 +0100
Subject: [PATCH 04/12] refactor: restructure methods

---
 src/qemcmc/sampler/qe_proposal.py | 37 ++++++++-----------------------
 1 file changed, 9 insertions(+), 28 deletions(-)

diff --git a/src/qemcmc/sampler/qe_proposal.py b/src/qemcmc/sampler/qe_proposal.py
index 6d105cb..89b10c1 100644
--- a/src/qemcmc/sampler/qe_proposal.py
+++ b/src/qemcmc/sampler/qe_proposal.py
@@ -1,6 +1,5 @@
 import warnings
 from typing import Optional
-
 import numpy as np
 
 from qemcmc.circuits import CircuitMaker
@@ -17,7 +16,7 @@ class QeProposal(Proposal):
 
     This class implements the proposal mechanism of the quantum-enhanced
     MCMC algorithm. Candidate states are generated by simulating the
-    quantum time evolution of a transverse-field Hamiltonian. 
+    quantum time evolution of a transverse-field Hamiltonian.
 
     The quantum proposal circuit is constructed using :class:`CircuitMaker`
     and can optionally operate on coarse-grained subgroups of spins to
@@ -43,7 +42,7 @@ class QeProposal(Proposal):
         sampled uniformly from the range. This allows for dynamic adjustment
         of the influence of different terms in the Hamiltonian during the proposal
         step. By default None.
-        Note that this is further adjusted by (1 - gamma) to balance the influence of the 
+        Note that this is further adjusted by (1 - gamma) to balance the influence of the
         problem Hamiltonian with the mixer term. Divide by (1 - gamma) if needed.
     m : int, optional
         Number of subgroups to partition the spins into for sequential updates.
@@ -72,8 +71,7 @@ def __init__(
         if coupling_weights is not None:
             if len(coupling_weights) != len(model.normalised_couplings):
                 raise ValueError(
-                    f"Length of coupling_weights ({len(coupling_weights)}) must match the number of "
-                    f"couplings in the model ({len(model.normalised_couplings)}), including constraint terms."
+                    f"Length of coupling_weights ({len(coupling_weights)}) must match the number of couplings in the model ({len(model.normalised_couplings)}), including constraint terms."
                 )
             self.coupling_weights = coupling_weights
         else:
@@ -87,7 +85,6 @@ def __init__(
 
         self.CM = CircuitMaker(self.model, delta_time=self.delta_time)
 
-        
         self.cg = coarse_graining or CoarseGraining(model.n_spins)
 
     def update(self, current_state: str) -> str:
@@ -108,33 +105,18 @@ def update(self, current_state: str) -> str:
         if not isinstance(current_state, str):
             raise TypeError(f"Bitstring must be of type str, got {type(current_state)}: {current_state!r}")
 
-        
-
         # Generate m disjoint partitions of the spins
         partitions = self.cg.get_partitions(m=self.m)
-        
-        final_coupling_weights, g, t = self.sample_hyperparams()
-
-
 
+        final_coupling_weights, g, t = self.sample_hyperparams()
 
         working_state = current_state
         for subgroup in partitions:
             # Recalculate local couplings as spins outside the subgroup may have flipped
-            local_couplings = self.model.get_subgroup_couplings(
-                subgroup=subgroup,
-                current_state=working_state,
-                coupling_weights=final_coupling_weights
-            )
-            working_state = self.CM.update(
-                s=working_state,
-                subgroup_choice=subgroup,
-                local_couplings=local_couplings,
-                gamma=g,
-                time=t
-            )
+            local_couplings = self.model.get_subgroup_couplings(subgroup=subgroup, current_state=working_state, coupling_weights=final_coupling_weights)
+            working_state = self.CM.update(s=working_state, subgroup_choice=subgroup, local_couplings=local_couplings, gamma=g, time=t)
         return working_state
-    
+
     def sample_hyperparams(self) -> tuple[list[float], float, int]:
         """
         Sample coupling weights for the current proposal step based on the provided coupling_weights configuration.
@@ -146,7 +128,7 @@ def sample_hyperparams(self) -> tuple[list[float], float, int]:
         # Pre-sample weights for all unique coupling weight settings (tuples or floats)
         unique_items = sorted(list(set(self.coupling_weights)), key=lambda x: str(x))
         unique_index = [[i for i, x in enumerate(self.coupling_weights) if x == item] for item in unique_items]
-        
+
         sampled_weights = np.ones(len(unique_items), dtype=float)
         for i, item in enumerate(unique_items):
             if isinstance(item, (int, float)):
@@ -155,7 +137,7 @@ def sample_hyperparams(self) -> tuple[list[float], float, int]:
                 sampled_weights[i] = np.random.uniform(*item)
 
         # Adjust the sampled weights by (1 - gamma) to balance their influence with the mixer term in the Hamiltonian
-        sampled_weights_adjusted_by_gamma = (1-g) * sampled_weights  # Adjust weights by (1 - gamma) to balance with the mixer term
+        sampled_weights_adjusted_by_gamma = (1 - g) * sampled_weights  # Adjust weights by (1 - gamma) to balance with the mixer term
 
         # Construct the final list of coupling weights for this update step
         final_coupling_weights = np.ones(len(self.coupling_weights))
@@ -163,7 +145,6 @@ def sample_hyperparams(self) -> tuple[list[float], float, int]:
             for index in indices:
                 final_coupling_weights[index] = sampled_weights_adjusted_by_gamma[i]
 
-        
         return final_coupling_weights, g, t
 
     def _validate_gamma(self, gamma: float | tuple[float, float]) -> float | tuple[float, float]:

From 9f366985593ae863da423c4ac074c70c6d6c99d1 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 16:58:14 +0100
Subject: [PATCH 05/12] refactor: re structure code

---
 src/qemcmc/sampler/runners.py | 51 ++++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/src/qemcmc/sampler/runners.py b/src/qemcmc/sampler/runners.py
index 5bd5d1f..cb40926 100644
--- a/src/qemcmc/sampler/runners.py
+++ b/src/qemcmc/sampler/runners.py
@@ -1,23 +1,25 @@
 import numpy as np
 from tqdm.auto import tqdm
-from typing import Optional, Callable, Tuple
+from typing import Optional, Tuple
 from qemcmc.model import EnergyModel, ConstraintModel
 from qemcmc.utils import MCMCChain, MCMCState, get_random_state
 from qemcmc.sampler.proposal import Proposal
 
+
 class Runner:
     """
-    Base class for running MCMC routines. 
+    Base class for running MCMC routines.
     Subclasses implement specific MCMC based algorithms.
     """
+
     def __init__(self):
         pass
-    
+
     def test_probs(self, energy_s: float, energy_sprime: float, temperature: float = 1.0) -> float:
         """
         Calculate the Metropolis acceptance probability.
 
-        This computes exp(-(E(s') - E(s)) / T), used to determine the acceptance 
+        This computes exp(-(E(s') - E(s)) / T), used to determine the acceptance
         probability of a new state s' given the current state s.
 
         Parameters
@@ -39,7 +41,7 @@ def test_probs(self, energy_s: float, energy_sprime: float, temperature: float =
             return 1.0
         else:
             exp_factor = np.exp(-delta_energy / temperature)
-        
+
         return min(1.0, exp_factor)
 
     def test_accept(self, energy_s: float, energy_sprime: float, temperature: float = 1.0) -> bool:
@@ -81,12 +83,21 @@ class MCMCRunner(Runner):
     temp : float
         The temperature for the Metropolis acceptance criterion.
     """
+
     def __init__(self, model: EnergyModel, temp: float):
         super().__init__()
         self.model = model
         self.temp = temp
 
-    def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = None, name: Optional[str] = None, verbose: bool = False, sample_frequency: int = 1) -> MCMCChain:
+    def run(
+        self,
+        proposer: Proposal,
+        n_hops: int,
+        initial_state: Optional[str] = None,
+        name: Optional[str] = None,
+        verbose: bool = False,
+        sample_frequency: int = 1,
+    ) -> MCMCChain:
         """
         Run the MCMC simulation.
 
@@ -136,6 +147,7 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
                 energy_s = energy_sprime
                 current_state = MCMCState(s_prime, accepted, energy_s, position=i)
 
+            # TODO: What is this doing? Why are we ignoring the step at i=0?
             if i % sample_frequency == 0 and i != 0:
                 mcmc_chain.add_state(MCMCState(current_state.bitstring, True, energy_s, position=i))
 
@@ -158,6 +170,7 @@ class ConstrainedMCMCRunner(Runner):
     reject_invalid : bool, optional
         If True (default), proposed states that violate the constraint are rejected.
     """
+
     def __init__(self, model: ConstraintModel, temp: float, reject_invalid: bool = True):
         if not isinstance(model, ConstraintModel):
             raise TypeError("Model must be an instance of ConstraintModel.")
@@ -168,7 +181,16 @@ def __init__(self, model: ConstraintModel, temp: float, reject_invalid: bool = T
         self.constraint_func = self.model.constraint_func
         self.reject_invalid = reject_invalid
 
-    def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = None, name: Optional[str] = None, verbose: bool = False, sample_frequency: int = 1, return_rejections: bool = True) -> Tuple[MCMCChain, int]:
+    def run(
+        self,
+        proposer: Proposal,
+        n_hops: int,
+        initial_state: Optional[str] = None,
+        name: Optional[str] = None,
+        verbose: bool = False,
+        sample_frequency: int = 1,
+        return_rejections: bool = True,
+    ) -> Tuple[MCMCChain, int]:
         """
         Run the constrained MCMC simulation.
 
@@ -195,6 +217,7 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
         tuple[MCMCChain, int]
             A tuple containing the generated Markov chain and the number of rejections due to constraints. If return_rejections is False, only the MCMCChain is returned.
         """
+
         if name is None:
             name = getattr(proposer, "method", "Constrained") + " MCMC"
 
@@ -208,7 +231,7 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
                     break
             if initial_state is None:
                 raise ValueError("Could not find a valid initial state. Please provide one.")
-        
+
         elif not self.constraint_func(initial_state):
             raise ValueError(f"Provided initial state {initial_state} does not satisfy the constraint.")
 
@@ -223,11 +246,13 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
         MH_rejects = 0
         s_prime = None
         pbar = tqdm(range(0, n_hops), desc="Run " + name, disable=not verbose)
-        E_diffs = [] # energy difference in proposal
-        h_diffs = [] # Hamming distance difference in proposal
+        E_diffs = []  # energy difference in proposal
+        h_diffs = []  # Hamming distance difference in proposal
         for i in pbar:
             s_prime = proposer.update(current_state.bitstring)
-            pbar.set_description(f"Run {name} | current state: {current_state.bitstring} | proposing: {s_prime} | avgEdiff: {np.mean(np.abs(E_diffs)):.4f} | avgHdiff: {np.mean(h_diffs):.2f} | constrejecects: {constraint_rejections} | selfrejects: {self_rejections} | MHrejects: {MH_rejects}")
+            pbar.set_description(
+                f"Run {name} | current state: {current_state.bitstring} | proposing: {s_prime} | avgEdiff: {np.mean(np.abs(E_diffs)):.4f} | avgHdiff: {np.mean(h_diffs):.2f} | constrejecects: {constraint_rejections} | selfrejects: {self_rejections} | MHrejects: {MH_rejects}"
+            )
             if s_prime == current_state.bitstring:
                 accepted = False
                 energy_sprime = energy_s
@@ -245,7 +270,8 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
             if accepted:
                 energy_s = energy_sprime
                 current_state = MCMCState(s_prime, accepted, energy_s, position=i)
-                
+
+            # TODO: What is this doing? Why are we ignoring the step at i=0?
             if i % sample_frequency == 0 and i != 0:
                 mcmc_chain.add_state(MCMCState(current_state.bitstring, True, energy_s, position=i))
 
@@ -253,4 +279,3 @@ def run(self, proposer: Proposal, n_hops: int, initial_state: Optional[str] = No
             return mcmc_chain, constraint_rejections
         else:
             return mcmc_chain
-    

From 05b1f20d288df0db94436ec49d32af689776b13d Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 17:00:45 +0100
Subject: [PATCH 06/12] fix: normalisation calculation

---
 src/qemcmc/model/energy_model.py | 64 +++++++++++++++++---------------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/src/qemcmc/model/energy_model.py b/src/qemcmc/model/energy_model.py
index 4ee2f13..3a2a976 100644
--- a/src/qemcmc/model/energy_model.py
+++ b/src/qemcmc/model/energy_model.py
@@ -6,7 +6,7 @@
 import math
 from tqdm import tqdm
 from qemcmc.utils.helpers import get_random_state
-import warnings
+
 
 class EnergyModel:
     """
@@ -37,20 +37,13 @@ class EnergyModel:
       small systems.
     """
 
-    def __init__(
-        self,
-        n: int,
-        couplings: List[np.ndarray] = [],
-        name: str = None,
-        cost_function_signs: list = [-1, -1],
-        model_type: str = "ising"
-    ):
+    def __init__(self, n: int, couplings: List[np.ndarray] = [], name: str = None, cost_function_signs: list = [-1, -1], model_type: str = "ising"):
         self.n = n
         self.n_spins = n
         self.couplings = couplings
         self.name = name
         self.alphas = self.calculate_alpha(n, couplings)
-
+        self.lowest_energy = None
 
         self.normalised_couplings = [self.couplings[i] * self.alphas[i] for i in range(len(self.couplings))]
         self.cost_function_signs = cost_function_signs
@@ -62,10 +55,10 @@ def __init__(
 
         self.initial_state = self.get_initial_states(num_initial_states=100)
 
-    def get_initial_states(self, num_initial_states:int):
+    def get_initial_states(self, num_initial_states: int):
         """
         Generates a list of random initial states.
-        
+
         Parameters
         ----------
         num_initial_states:
@@ -114,7 +107,7 @@ def calculate_energy(self, state, couplings, cost_function_signs):
         """
         Calculate the energy of a given state for an arbitrary-order Ising/binary model.
 
-        Parameters:
+        Parameters
         -----------
             state : array-like (str, list, tuple, np.array)
                 State configuration. Can be:
@@ -127,7 +120,7 @@ def calculate_energy(self, state, couplings, cost_function_signs):
                 - 2D arrays represent quadratic terms (J_ij)
                 - 3D arrays represent cubic terms, etc.
 
-        Returns:
+        Returns
         --------
             float : Total energy of the state
         """
@@ -158,7 +151,7 @@ def calculate_energy(self, state, couplings, cost_function_signs):
                 else:
                     # Convert to spin
                     state = np.array([2 * int(bit) - 1 for bit in state])
-        
+
         total_energy = 0.0
         for term_index, coupling in enumerate(couplings):
             coupling = np.array(coupling)
@@ -205,14 +198,13 @@ def get_subgroup_couplings(self, subgroup: List[int], current_state: str, coupli
 
         Notes
         -----
-                It might seem off to do the reweighting at this point, but here are reasons why I [SF] did it! 
-                Basically, when you get the subgroup couplings, the terms jumble up, so the returned local 
+                It might seem off to do the reweighting at this point, but here are reasons why I [SF] did it!
+                Basically, when you get the subgroup couplings, the terms jumble up, so the returned local
                 coupling list necessarily does not respect the order etc. of the different terms in the original.
         """
 
-        
         if coupling_weights is None:
-            coupling_weights = [1.0] * len(subgroup)
+            coupling_weights = [1.0] * len(self.normalised_couplings)
 
         n_sub = len(subgroup)
         subgroup_set = set(subgroup)
@@ -227,7 +219,7 @@ def get_subgroup_couplings(self, subgroup: List[int], current_state: str, coupli
             coupling = np.array(coupling) * coupling_weights[couplings_index]
             # only loop over elements that actually exist (non-zero)
             non_zero_indices = np.transpose(np.nonzero(coupling))
-            
+
             for indices in non_zero_indices:
                 indices = tuple(indices)
                 coeff = coupling[indices]
@@ -246,6 +238,7 @@ def get_subgroup_couplings(self, subgroup: List[int], current_state: str, coupli
                     new_order = len(in_group)
                     local_indices = tuple(g_to_l[i] for i in in_group)
                     new_couplings[new_order - 1][local_indices] += effective_coeff
+
         return new_couplings
 
     def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) -> np.ndarray:
@@ -266,13 +259,14 @@ def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) ->
 
         Returns
         -------
-        np.ndarray : 
+        np.ndarray :
             An array of normalising factors for each term in the couplings list.
         """
         if couplings is None:
             couplings = self.couplings
 
         norm_sq_arr = np.zeros(len(couplings))
+
         for T_ind, T in enumerate(couplings):
             norm_sq = 0.0
             T = np.asarray(T)
@@ -282,7 +276,7 @@ def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) ->
                 pass
 
             # 1-body: h_i
-            if order == 1:
+            elif order == 1:
                 if T.shape != (n,):
                     raise ValueError(f"1-body tensor has shape {T.shape}, expected ({n},)")
                 for i in range(n):
@@ -290,11 +284,11 @@ def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) ->
                     norm_sq += c * c
 
             # 2-body: symmetric J
-            if order == 2:
+            elif order == 2:
                 if T.shape != (n, n):
                     raise ValueError(f"2-body tensor has shape {T.shape}, expected ({n},{n})")
 
-                # Enforce symmetry (rejects pure upper/lower triangular)
+                # Enforce symmetry (reject pure upper/lower triangular)
                 if not np.allclose(T, T.T):
                     raise ValueError("Non-symmetric J provided. This alpha function only accepts symmetric J.")
 
@@ -304,11 +298,11 @@ def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) ->
                         c = float(T[i, j])
                         if c != 0.0:
                             norm_sq += c * c
-                
 
             # Order >= 3: count each unordered interaction once using i1<i2<...<ik
-            if T.shape != (n,) * order:
-                raise ValueError(f"{order}-body tensor has shape {T.shape}, expected {(n,) * order}")
+            else:
+                if T.shape != (n,) * order:
+                    raise ValueError(f"{order}-body tensor has shape {T.shape}, expected {(n,) * order}")
 
             for comb in itertools.combinations(range(n), order):
                 c = float(T[comb])
@@ -320,10 +314,21 @@ def calculate_alpha(self, n: int, couplings: list = None, eps: float = 1e-15) ->
         if norm_sq_tot < eps:
             print("Warning: Couplings are all zero (or very close to zero). Returning alpha=1 to avoid division by zero.")
             print("Input of a zero-coupling may be intended to sample Uniform distributions, however if this is not your intention, please check your couplings.")
-            #return np.ones(len(couplings))
+            return np.ones(len(couplings))
 
         return np.sqrt(n / norm_sq_arr)
 
+        # TODO: is this check required? test code with and without.
+        # alpha_arr = np.zeros(len(couplings))
+        # for i, val in enumerate(norm_sq_arr):
+        #     if val < eps:
+        #         print(f"Warning: Coupling at index {i} is all zero (or close to zero). Returning alpha=1 for this term.")
+        #         alpha_arr[i] = 1.0
+        #     else:
+        #         alpha_arr[i] = np.sqrt(n / val)
+
+        # return alpha_arr
+
     def get_energy(self, state: str) -> float:
         """
         Returns the energy of a given state
@@ -425,7 +430,7 @@ def get_lowest_energy(self):
             all_energies = self.get_all_energies()
 
         lowest_energy = np.min(all_energies)
-
+        self.lowest_energy = lowest_energy
         return lowest_energy
 
     def get_boltzmann_factor(self, state: str, beta: float = 1.0) -> float:
@@ -457,4 +462,3 @@ def get_boltzmann_factor_from_energy(self, E, beta: float = 1.0) -> float:
         """
 
         return np.exp(-1 * beta * E, dtype=np.longdouble)
-    
\ No newline at end of file

From d31e79898e1275871066c7b31b8001def85c5f3d Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 17:02:05 +0100
Subject: [PATCH 07/12] refactor: minor changes

---
 src/qemcmc/circuits.py | 65 +++++++++++++++++++-----------------------
 1 file changed, 29 insertions(+), 36 deletions(-)

diff --git a/src/qemcmc/circuits.py b/src/qemcmc/circuits.py
index e8372bf..277bd14 100644
--- a/src/qemcmc/circuits.py
+++ b/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,12 +39,9 @@ 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.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."""
@@ -95,7 +93,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
@@ -114,9 +111,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
@@ -155,45 +152,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)
@@ -234,30 +227,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):
@@ -296,9 +289,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)

From c8ad525f0a0ac54a0186cce82db9dcb781100414 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 17:04:20 +0100
Subject: [PATCH 08/12] fix: return a list type in else block of get_partitions

---
 src/qemcmc/coarse_grain.py | 45 +++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/src/qemcmc/coarse_grain.py b/src/qemcmc/coarse_grain.py
index 09acd96..e418b5b 100644
--- a/src/qemcmc/coarse_grain.py
+++ b/src/qemcmc/coarse_grain.py
@@ -1,30 +1,24 @@
-# 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):
 
         if subgroups is None:
             subgroups = [list(range(n))]
@@ -40,8 +34,10 @@ 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):
+        """
+        Randomly samples a subgroup according to the specified probability distribution.
+        """
         idx = rng.choice(len(self.subgroups), p=self.subgroup_probs)
         return self.subgroups[idx]
 
@@ -57,5 +53,4 @@ def get_partitions(self, m: int, rng=np.random):
         if self.repeated:
             return chunks
         else:
-            return chunks[0]
-
+            return [chunks[0]]

From 3da6796f0cf81d7d343d09d173a6428dafdd79f0 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 17:05:05 +0100
Subject: [PATCH 09/12] fix: comment old code for now

---
 src/qemcmc/main.py | 246 ++++++++++++++++++++++-----------------------
 1 file changed, 123 insertions(+), 123 deletions(-)

diff --git a/src/qemcmc/main.py b/src/qemcmc/main.py
index 6b5877c..226f64e 100644
--- a/src/qemcmc/main.py
+++ b/src/qemcmc/main.py
@@ -1,123 +1,123 @@
-# Internal
-from qemcmc.sampler import ClassicalMCMC, QeMCMC
-from qemcmc.sampler.runners import MCMCRunner
-from qemcmc.utils import ModelMaker, MCMCChain
-from qemcmc.coarse_grain import CoarseGraining
-
-
-# External
-from matplotlib import pyplot as plt
-from joblib import Parallel, delayed
-import numpy as np
-from tqdm import tqdm
-
-
-def plot_thermalisation(chains_dict, n_spins, temp, exact_energy=None, lowest_levels=None):
-    """
-    chains_dict format: {"Label": (chain_data, "color")}
-    """
-    fig, ax = plt.subplots(figsize=(8, 6))
-
-    # Plot each algorithm's energy trajectory
-    for label, (energy_data, color) in chains_dict.items():
-        steps = np.arange(1, len(energy_data) + 1)
-        ax.plot(steps, energy_data, label=label, color=color, linewidth=2)
-
-    # Add reference lines if provided
-    if exact_energy is not None:
-        ax.axhline(exact_energy, color="black", label="Exact average energy", zorder=3)
-
-    if lowest_levels is not None:
-        for i, level in enumerate(lowest_levels):
-            lbl = "Lowest energy levels" if i == 0 else ""
-            ax.axhline(level, color="lightcyan", linestyle="--", label=lbl, zorder=1)
-
-    # Match the paper's style
-    ax.set_xscale("log")
-    ax.set_xlabel("MCMC step (Log Scale)", fontsize=12)
-    ax.set_ylabel("Average Energy", fontsize=12)
-    ax.set_title(f"Repeated Coarse Graining Benchmarks ({n_spins} spins | T = {temp})", fontsize=14)
-    ax.legend(loc="upper right", fontsize=9)
-
-    plt.tight_layout()
-    plt.show()
-
-
-if __name__ == "__main__":
-    n = 24
-    steps = 10
-    reps = 10
-    temp = 0.1
-
-    model_maker = ModelMaker(n, "Coarse Grained Ising", "24 Spin Test")
-    model = model_maker.model
-    initial_states = model.initial_state
-
-    cg = CoarseGraining(n)
-
-    m_values = [3, 4, 6, 8, 12]
-    colors = ["purple", "red", "pink", "green", "orange"]
-
-    results_dict = {}
-    runner = MCMCRunner()
-
-    # --- 1. Run Classical Baselines ---
-    print("--- Running Classical Baselines ---")
-
-    # Uniform
-    def run_uniform(rep):
-        sampler = ClassicalMCMC(model, temp, method="uniform")
-        return runner.run(sampler, steps, initial_state=initial_states[rep], verbose=True)
-
-    uni_chains = list(tqdm(Parallel(n_jobs=-1, return_as="generator")(delayed(run_uniform)(rep) for rep in range(reps)), total=reps, desc="Running Uniform chains", leave=False))
-    results_dict["Uniform"] = (np.mean([c.get_current_energy_array() for c in uni_chains], axis=0), "darkorange")
-
-    # Local
-    def run_local(rep):
-        sampler = ClassicalMCMC(model, temp, method="local")
-        return runner.run(sampler, steps, initial_state=initial_states[rep], verbose=True)
-
-    loc_chains = list(tqdm(Parallel(n_jobs=-1, return_as="generator")(delayed(run_local)(rep) for rep in range(reps)), total=reps, desc="Running Local chains", leave=False))
-    results_dict["Local"] = (np.mean([c.get_current_energy_array() for c in loc_chains], axis=0), "forestgreen")
-
-    # --- 2. Run Repeated QeMCMC ---
-    print("\n--- Running Quantum Chains (Repeated CG) ---")
-    m_values = [3, 4, 6, 8, 12]
-    colors = ["purple", "red", "pink", "dodgerblue", "gold"]
-
-    for m, color in tqdm(zip(m_values, colors), total=len(m_values), desc="Testing Partitions (m)"):
-
-        def run_repeated_qemcmc(rep):
-            sampler = QeMCMC(
-                model=model,
-                gamma=(0.3, 0.6),
-                time=(2, 20),
-                temp=temp,
-                coarse_graining=cg,
-                m=m,
-            )
-            return runner.run(
-                sampler,
-                steps,
-                initial_state=initial_states[rep],
-                name=f"QeMCMC m={m}",
-                verbose=True,
-                sample_frequency=1,
-            )
-
-        chains = list(
-            tqdm(
-                Parallel(n_jobs=-1, return_as="generator")(delayed(run_repeated_qemcmc)(rep) for rep in range(reps)),
-                total=reps,
-                desc=f"Running m={m} chains",
-                leave=True,
-            )
-        )
-
-        all_energies = np.array([chain.get_current_energy_array() for chain in chains])
-        avg_energy = np.mean(all_energies, axis=0)
-
-        results_dict[f"QeMCMC m={m} ({n // m} spin blocks)"] = (avg_energy, color)
-
-    exact_gs_energy = model.get_ground_state()
-    plot_thermalisation(results_dict, n_spins=n, temp=temp, exact_energy=exact_gs_energy)
+# # Internal
+# from qemcmc.sampler import ClassicalMCMC, QeMCMC
+# from qemcmc.sampler.runners import MCMCRunner
+# from qemcmc.utils import ModelMaker, MCMCChain
+# from qemcmc.coarse_grain import CoarseGraining
+
+
+# # External
+# from matplotlib import pyplot as plt
+# from joblib import Parallel, delayed
+# import numpy as np
+# from tqdm import tqdm
+
+
+# def plot_thermalisation(chains_dict, n_spins, temp, exact_energy=None, lowest_levels=None):
+#     """
+#     chains_dict format: {"Label": (chain_data, "color")}
+#     """
+#     fig, ax = plt.subplots(figsize=(8, 6))
+
+#     # Plot each algorithm's energy trajectory
+#     for label, (energy_data, color) in chains_dict.items():
+#         steps = np.arange(1, len(energy_data) + 1)
+#         ax.plot(steps, energy_data, label=label, color=color, linewidth=2)
+
+#     # Add reference lines if provided
+#     if exact_energy is not None:
+#         ax.axhline(exact_energy, color="black", label="Exact average energy", zorder=3)
+
+#     if lowest_levels is not None:
+#         for i, level in enumerate(lowest_levels):
+#             lbl = "Lowest energy levels" if i == 0 else ""
+#             ax.axhline(level, color="lightcyan", linestyle="--", label=lbl, zorder=1)
+
+#     # Match the paper's style
+#     ax.set_xscale("log")
+#     ax.set_xlabel("MCMC step (Log Scale)", fontsize=12)
+#     ax.set_ylabel("Average Energy", fontsize=12)
+#     ax.set_title(f"Repeated Coarse Graining Benchmarks ({n_spins} spins | T = {temp})", fontsize=14)
+#     ax.legend(loc="upper right", fontsize=9)
+
+#     plt.tight_layout()
+#     plt.show()
+
+
+# if __name__ == "__main__":
+#     n = 24
+#     steps = 10
+#     reps = 10
+#     temp = 0.1
+
+#     model_maker = ModelMaker(n, "Coarse Grained Ising", "24 Spin Test")
+#     model = model_maker.model
+#     initial_states = model.initial_state
+
+#     cg = CoarseGraining(n)
+
+#     m_values = [3, 4, 6, 8, 12]
+#     colors = ["purple", "red", "pink", "green", "orange"]
+
+#     results_dict = {}
+#     runner = MCMCRunner()
+
+#     # --- 1. Run Classical Baselines ---
+#     print("--- Running Classical Baselines ---")
+
+#     # Uniform
+#     def run_uniform(rep):
+#         sampler = ClassicalMCMC(model, temp, method="uniform")
+#         return runner.run(sampler, steps, initial_state=initial_states[rep], verbose=True)
+
+#     uni_chains = list(tqdm(Parallel(n_jobs=-1, return_as="generator")(delayed(run_uniform)(rep) for rep in range(reps)), total=reps, desc="Running Uniform chains", leave=False))
+#     results_dict["Uniform"] = (np.mean([c.get_current_energy_array() for c in uni_chains], axis=0), "darkorange")
+
+#     # Local
+#     def run_local(rep):
+#         sampler = ClassicalMCMC(model, temp, method="local")
+#         return runner.run(sampler, steps, initial_state=initial_states[rep], verbose=True)
+
+#     loc_chains = list(tqdm(Parallel(n_jobs=-1, return_as="generator")(delayed(run_local)(rep) for rep in range(reps)), total=reps, desc="Running Local chains", leave=False))
+#     results_dict["Local"] = (np.mean([c.get_current_energy_array() for c in loc_chains], axis=0), "forestgreen")
+
+#     # --- 2. Run Repeated QeMCMC ---
+#     print("\n--- Running Quantum Chains (Repeated CG) ---")
+#     m_values = [3, 4, 6, 8, 12]
+#     colors = ["purple", "red", "pink", "dodgerblue", "gold"]
+
+#     for m, color in tqdm(zip(m_values, colors), total=len(m_values), desc="Testing Partitions (m)"):
+
+#         def run_repeated_qemcmc(rep):
+#             sampler = QeMCMC(
+#                 model=model,
+#                 gamma=(0.3, 0.6),
+#                 time=(2, 20),
+#                 temp=temp,
+#                 coarse_graining=cg,
+#                 m=m,
+#             )
+#             return runner.run(
+#                 sampler,
+#                 steps,
+#                 initial_state=initial_states[rep],
+#                 name=f"QeMCMC m={m}",
+#                 verbose=True,
+#                 sample_frequency=1,
+#             )
+
+#         chains = list(
+#             tqdm(
+#                 Parallel(n_jobs=-1, return_as="generator")(delayed(run_repeated_qemcmc)(rep) for rep in range(reps)),
+#                 total=reps,
+#                 desc=f"Running m={m} chains",
+#                 leave=True,
+#             )
+#         )
+
+#         all_energies = np.array([chain.get_current_energy_array() for chain in chains])
+#         avg_energy = np.mean(all_energies, axis=0)
+
+#         results_dict[f"QeMCMC m={m} ({n // m} spin blocks)"] = (avg_energy, color)
+
+#     exact_gs_energy = model.get_ground_state()
+#     plot_thermalisation(results_dict, n_spins=n, temp=temp, exact_energy=exact_gs_energy)

From 4b289ccfb3b870a15cd1627b81d721521a7c577f Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 18:43:22 +0100
Subject: [PATCH 10/12] feat: add update method

---
 src/qemcmc/sampler/classical_proposal.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/qemcmc/sampler/classical_proposal.py b/src/qemcmc/sampler/classical_proposal.py
index b10589f..2a7a1ee 100644
--- a/src/qemcmc/sampler/classical_proposal.py
+++ b/src/qemcmc/sampler/classical_proposal.py
@@ -27,18 +27,21 @@ class ClassicalProposal(Proposal):
         Default is ``"uniform"``.
     """
 
+    METHODS = {"uniform", "local", "2-local"}
+
     def __init__(self, model: EnergyModel, method: str = "uniform"):
+        if method not in self.METHODS:
+            raise ValueError(f"Method '{method}' is not supported. Choose from 'uniform', 'local', or '2-local'.")
         super().__init__(model)
         self.method = method
 
+    def update(self, current_state_bitstring: str) -> str:
         if self.method == "uniform":
-            self.update = self.update_uniform
+            return self.update_uniform(current_state_bitstring)
         elif self.method == "local":
-            self.update = self.update_local
-        elif self.method == "2-local":
-            self.update = self.update_2local
+            return self.update_local(current_state_bitstring)
         else:
-            raise ValueError(f"Method '{method}' is not supported. Choose from 'uniform', 'local', or '2-local'.")
+            return self.update_2local(current_state_bitstring)
 
     def update_uniform(self, current_state_bitstring: str) -> str:
         """

From 968718daab7ca213b457a8de1c54ba778c30ee72 Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 18:46:35 +0100
Subject: [PATCH 11/12] docs: add params, return typeS

---
 src/qemcmc/model/energy_model.py | 37 ++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/qemcmc/model/energy_model.py b/src/qemcmc/model/energy_model.py
index 3a2a976..9cce5f7 100644
--- a/src/qemcmc/model/energy_model.py
+++ b/src/qemcmc/model/energy_model.py
@@ -358,16 +358,19 @@ def get_lowest_energies(self, num_states: int, return_configurations: bool = Fal
         """
         Retrieve the lowest energy states and their degeneracies.
         This method computes all possible energies and then finds the specified number
-        of lowest energy states along with their degeneracies. Note that this method
-        is intended for small instances due to its brute-force nature, which is extremely
-        memory intensive and slow.
-        Args:
+        of lowest energy states along with their degeneracies.
+
+        Parameters
             num_states (int): The number of lowest energy states to retrieve.
             return_configurations (bool): Whether to also return the corresponding configurations of the lowest energy states. Defaults to False.
-        Returns:
+
+        Returns
             Two numpy arrays:
                 - The first array contains the lowest energy values.
                 - The second array contains the degeneracies of the corresponding energy values.
+
+        Notes
+            This method is intended for small instances due to its brute-force nature, which is extremely memory intensive and slow.
         """
         # only to be used for small instances, it is just brute force so extremely memory intensive and slow
         all_energies = self.get_all_energies()
@@ -389,14 +392,14 @@ def find_lowest_values(self, arr: np.ndarray, num_values: int = 5):
         """
         Find the lowest unique values in an array and their degeneracies.
 
-        Args:
+        Parameters
             arr (np.ndarray): The input array from which to find the lowest values.
             num_values (int, optional): The number of lowest unique values to find. Defaults to 5.
 
-        Returns:
-        tuple: A tuple containing two numpy arrays:
-            - lowest_values (np.ndarray): The lowest unique values in the array.
-            - degeneracy (np.ndarray): The counts of each of the lowest unique values.
+        Returns
+            tuple: A tuple containing two numpy arrays:
+                - lowest_values (np.ndarray): The lowest unique values in the array.
+                - degeneracy (np.ndarray): The counts of each of the lowest unique values.
         """
         # Count the occurrences of each value
         unique_values, counts = np.unique(arr, return_counts=True)
@@ -415,9 +418,11 @@ def get_lowest_energy(self):
         This method uses a brute force approach to find the lowest energy,
         making it extremely memory intensive and slow. It is recommended
         to use this method only for small instances.
-        Returns:
+
+        Returns
             float: The lowest energy value.
-        Notes:
+
+        Notes
             If the lowest energy has already been calculated and stored
             in `self.lowest_energy`, it will return that value directly
             to save computation time.
@@ -437,11 +442,11 @@ def get_boltzmann_factor(self, state: str, beta: float = 1.0) -> float:
         """
         Get un-normalised boltzmann probability of a given state
 
-        Args:
+        Parameters
             state (str): configuration of spins for which probability is to be calculated
             beta (float): inverse temperature (1/T) at which the probability is to be calculated.
 
-        Returns:
+        Returns
             float corresponding to the un-normalised boltzmann probability of the given state.
         """
         E = self.get_energy(state)
@@ -453,11 +458,11 @@ def get_boltzmann_factor_from_energy(self, E, beta: float = 1.0) -> float:
         """
         Get un-normalized Boltzmann probability for a given energy.
 
-        Args:
+        Parameters
             E (float): Energy for which the Boltzmann factor is to be calculated.
             beta (float): Inverse temperature (1/T) at which the probability is to be calculated.
 
-        Returns:
+        Returns
             float: The un-normalized Boltzmann probability for the given energy.
         """
 

From 2997474a80d5febc243497d474869c38979cefec Mon Sep 17 00:00:00 2001
From: Feroz Hassan <75996466+ferozmay@users.noreply.github.com>
Date: Mon, 13 Apr 2026 18:47:33 +0100
Subject: [PATCH 12/12] fix: include user specified subgroups in get_partitions

---
 src/qemcmc/coarse_grain.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/qemcmc/coarse_grain.py b/src/qemcmc/coarse_grain.py
index e418b5b..64e8fd1 100644
--- a/src/qemcmc/coarse_grain.py
+++ b/src/qemcmc/coarse_grain.py
@@ -20,6 +20,8 @@ class CoarseGraining:
 
     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))]
             subgroup_probs = [1.0]
@@ -34,20 +36,31 @@ 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):
+    def sample(self) -> list[int]:
         """
         Randomly samples a subgroup according to the specified probability distribution.
         """
-        idx = rng.choice(len(self.subgroups), p=self.subgroup_probs)
+        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: