Add x0 functionality to mcmc sampling.

Author: arm61

kinisi/fitting.py

@@ -201,17 +201,23 @@ def max_aposteriori(self) -> None:
         for i, name in enumerate(self.parameter_names):
             self.data_group[name] = result[i] * self.parameter_units[i]
 
-    def mcmc(self, n_samples: int = 1000, n_walkers: int = 32, n_burn: int = 500, n_thin=10) -> None:
+    def mcmc(self, x0: tuple[sc.Variable] = None, n_samples: int = 1000, n_walkers: int = 32, n_burn: int = 500, n_thin=10) -> None:
         """
         Perform MCMC sampling of the model parameters.
 
-        :param n_samples: Number of samples to generate
-        :param n_walkers: Number of MCMC walkers
-        :param n_burn: Number of burn-in samples
-        :param n_thin: Thinning factor
+        :param x0: Initial starting position for MCMC sampling. Optional, defaults to max likelihood/aposteriori or mean of samples.
+        :param n_samples: Number of samples to generate. Optional, defaults to 1000. 
+        :param n_walkers: Number of MCMC walkers. Optional, defaults to 32.
+        :param n_burn: Number of burn-in samples. Optional, defaults to 500.
+        :param n_thin: Thinning factor. Optional, defaults to 10.
         """
         if isinstance(self.data_group[self.parameter_names[0]], Samples):
             values = np.array([sc.mean(self.data_group[p]).value for p in self.parameter_names])
+        elif x0 is not None: 
+            for i, x in enumerate(x0):
+                if self.data_group[self.parameter_names[i]].unit != x.unit:
+                    raise TypeError("x0 input must have same unit as paramters")
+            values = np.array([x.value for x in x0])
         else:
             values = np.array([self.data_group[p].value for p in self.parameter_names])
         pos = values + values * 1e-2 * np.random.randn(n_walkers, len(self.parameter_names))

kinisi/tests/test_fitting.py

@@ -9,6 +9,7 @@
 # pylint: disable=R0201
 
 import unittest
+import pytest
 
 import numpy as np
 import scipp as sc
@@ -142,6 +143,25 @@ def test_mcmc(self):
         assert isinstance(td.data_group['c'], Samples)
         assert td.data_group['m'].shape == (32,)
         assert td.data_group['c'].shape == (32,)
+    
+    def test_mcmc_x0(self):
+        """
+        Test the MCMC sampling function with an x0 value.
+        """
+        td = fitting.FittingBase(data, straight_line, ('m', 'c'), (sc.Unit('m/s'), sc.Unit('m')))
+        td.mcmc(x0=(1 * sc.Unit('m/s'), 0.5 * sc.Unit('m')), n_samples=10, n_burn=5, n_walkers=32)
+        assert isinstance(td.data_group['m'], Samples)
+        assert isinstance(td.data_group['c'], Samples)
+        assert td.data_group['m'].shape == (32,)
+        assert td.data_group['c'].shape == (32,) 
+
+    def test_mcmc_x0_wrong_unit(self):
+        """
+        Test the MCMC sampling function with an x0 value.
+        """
+        td = fitting.FittingBase(data, straight_line, ('m', 'c'), (sc.Unit('m/s'), sc.Unit('m')))
+        with pytest.raises(TypeError):
+            td.mcmc(x0=(1 * sc.Unit('m'), 0.5 * sc.Unit('m')), n_samples=10, n_burn=5, n_walkers=32)
 
     def test_nested_sampling(self):
         """