Fix docstring, default prior construction, return annotation, and test slope/viscosity labeling

Author: Copilot

kinisi/fitting.py

@@ -181,7 +181,7 @@ def prior_transform(self, parameters: tuple[float]) -> tuple[float]:
             x[i] = self.priors[i].ppf(parameters[i])
         return x
 
-    def max_likelihood(self):
+    def max_likelihood(self) -> None:
         """Find the max likelihood fit parameters for the model."""
         if self.priors is not None:
             x0 = [p.mean() for p in self.priors]
@@ -191,7 +191,7 @@ def max_likelihood(self):
         for i, name in enumerate(self.parameter_names):
             self.data_group[name] = result[i] * self.parameter_units[i]
 
-    def max_aposteriori(self):
+    def max_aposteriori(self) -> None:
         """Find the max aposteriori fit parameters for the model."""
         if self.priors is not None:
             x0 = [p.mean() for p in self.priors]

kinisi/tests/test_yeh_hummer.py

@@ -106,14 +106,14 @@ def test_yeh_hummer_priors(self):
         # Custom priors
         priors = (
             uniform(4e-5, 7e-5),  # D_0 prior
-            uniform(1e-4, 1e-2),  # viscosity prior
+            uniform(1e-4, 1e-2),  # slope prior
         )
 
         yh = YehHummer(td, temperature=sc.scalar(298, unit='K'), priors=priors)
 
         # Check that fitted values are within priors
         assert priors[0].a <= yh.D_infinite.value <= (priors[0].b + priors[0].a)
-        assert priors[1].a <= yh.shear_viscosity.value <= (priors[1].b + priors[1].a)
+        assert priors[1].a <= yh.data_group['slope'].value <= (priors[1].b + priors[1].a)
 
     def test_yeh_hummer_properties(self):
         """Test YehHummer property accessors."""

kinisi/yeh_hummer.py

@@ -41,8 +41,8 @@ class YehHummer(FittingBase):
 
     :param diffusion: sc.DataArray with diffusion coefficients and box_length coordinate
     :param temperature: Temperature (will be extracted from coords if not provided)
-    :param priors: Optional priors for [D_0, viscosity] parameters using scipy.stats objects
-        (viscosity in Pa*s)
+    :param priors: Optional priors for [D_0, slope] parameters using scipy.stats objects,
+        where slope has units of [diffusion] * [length]
     """
 
     def __init__(self, diffusion, temperature: sc.Variable, priors=None):
@@ -65,15 +65,14 @@ def __init__(self, diffusion, temperature: sc.Variable, priors=None):
         # Compute priors: use provided or defaults
         if priors is None:
             D_max = np.max(diffusion.values)
-            D_prior = uniform(D_max * 0.8, (D_max * 2.0) - (D_max - 0.8))
+            D_prior = uniform(D_max * 0.8, D_max * 1.2)
             visc_lower, visc_upper = 1e-5 * sc.Unit('Pa*s'), 1e-1 * sc.Unit('Pa*s')
 
-            # Higher viscosity = lower slope, so bounds are inverted
-            slope_bounds = (
-                self.viscosity_to_slope(visc_upper) * self._slope_unit,
-                self.viscosity_to_slope(visc_lower) * self._slope_unit,
-            )
-            slope_prior = uniform(slope_bounds[0], slope_bounds[1] - slope_bounds[0])
+            # Higher viscosity = lower slope, so bounds are inverted; viscosity_to_slope
+            # already returns plain floats so no unit multiplication needed here
+            slope_min = self.viscosity_to_slope(visc_upper)
+            slope_max = self.viscosity_to_slope(visc_lower)
+            slope_prior = uniform(slope_min, slope_max - slope_min)
             priors = [D_prior, slope_prior]
         else:
             if len(priors) != 2: