Fix coverage

Author: Tom-Willemsen

src/ibex_bluesky_core/devices/muon.py

@@ -36,7 +36,7 @@ def damped_oscillator(
     return B + A_0 * np.cos(omega_0 * t + phi_0) * np.exp(-t * lambda_0)
 
 
-def damped_oscillator_multiple(  # noqa: PLR0913 PLR0917 (model is just this complex)
+def double_damped_oscillator(  # noqa: PLR0913 PLR0917 (model is just this complex)
     t: NDArray[np.floating],
     B: float,  # noqa: N803
     A_0: float,  # noqa: N803

tests/devices/test_muon.py

@@ -1,5 +1,6 @@
 # pyright: reportMissingParameterType=false
-from unittest.mock import AsyncMock, patch
+import re
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import lmfit
 import numpy as np
@@ -10,7 +11,7 @@
 from ibex_bluesky_core.devices.muon import (
     MuonAsymmetryReducer,
     damped_oscillator,
-    damped_oscillator_multiple,
+    double_damped_oscillator,
 )
 from ibex_bluesky_core.devices.simpledae import MEventsWaiter, PeriodPerPointController, SimpleDae
 
@@ -23,19 +24,6 @@
 damped_oscillator_params.add("phi_0", 0.0)
 damped_oscillator_params.add("lambda_0", 0.001)
 
-multi_damped_oscillator_model = lmfit.Model(damped_oscillator_multiple)
-
-multi_damped_oscillator_params = lmfit.Parameters()
-multi_damped_oscillator_params.add("B", 0.0)
-multi_damped_oscillator_params.add("A_0", 0.1, min=0)
-multi_damped_oscillator_params.add("omega_0", 0.1, min=0)
-multi_damped_oscillator_params.add("phi_0", 0.0)
-multi_damped_oscillator_params.add("lambda_0", 0.001)
-multi_damped_oscillator_params.add("A_1", 0.1, min=0)
-multi_damped_oscillator_params.add("omega_1", 0.1, min=0)
-multi_damped_oscillator_params.add("phi_1", 0.0)
-multi_damped_oscillator_params.add("lambda_1", 0.001)
-
 
 @pytest.fixture
 def asymmetry_reducer():
@@ -61,18 +49,7 @@ def rebinning_asymmetry_reducer():
 
 
 @pytest.fixture
-def multi_damped_oscillator_asymmetry_reducer():
-    return MuonAsymmetryReducer(
-        forward_detectors=np.array([1]),
-        backward_detectors=np.array([2]),
-        prefix="UNITTEST:",
-        model=multi_damped_oscillator_model,
-        fit_parameters=multi_damped_oscillator_params,
-    )
-
-
-@pytest.fixture
-async def simpledae_single_oscillator(rebinning_asymmetry_reducer):
+async def simpledae(rebinning_asymmetry_reducer):
     dae = SimpleDae(
         prefix="UNITTEST:",
         reducer=rebinning_asymmetry_reducer,
@@ -83,63 +60,50 @@ async def simpledae_single_oscillator(rebinning_asymmetry_reducer):
     return dae
 
 
-@pytest.fixture
-async def simpledae_multi_oscillator(multi_damped_oscillator_asymmetry_reducer):
-    dae = SimpleDae(
-        prefix="UNITTEST:",
-        reducer=multi_damped_oscillator_asymmetry_reducer,
-        waiter=MEventsWaiter(5000),
-        controller=PeriodPerPointController(save_run=False),
+def test_damped_oscillator():
+    np.testing.assert_allclose(
+        damped_oscillator(
+            t=np.array([0, 1, 2, 3, 4, 5], dtype=np.float64),
+            B=1,
+            A_0=0.0,
+            phi_0=0.0,
+            omega_0=0.0,
+            lambda_0=0.0,
+        ),
+        np.array([1, 1, 1, 1, 1, 1], dtype=np.float64),
     )
-    await dae.connect(mock=True)
-    return dae
 
 
-def test_asymmetry_reducer_readable_signals_single_oscillator(
-    simpledae_single_oscillator, rebinning_asymmetry_reducer
-):
-    assert set(
-        rebinning_asymmetry_reducer.additional_readable_signals(simpledae_single_oscillator)
-    ) == {
-        simpledae_single_oscillator.reducer.B,
-        simpledae_single_oscillator.reducer.B_err,
-        simpledae_single_oscillator.reducer.A_0,
-        simpledae_single_oscillator.reducer.A_0_err,
-        simpledae_single_oscillator.reducer.omega_0,
-        simpledae_single_oscillator.reducer.omega_0_err,
-        simpledae_single_oscillator.reducer.phi_0,
-        simpledae_single_oscillator.reducer.phi_0_err,
-        simpledae_single_oscillator.reducer.lambda_0,
-        simpledae_single_oscillator.reducer.lambda_0_err,
-    }
+def test_double_damped_oscillator():
+    np.testing.assert_allclose(
+        double_damped_oscillator(
+            t=np.array([0, 1, 2, 3, 4, 5], dtype=np.float64),
+            B=1,
+            A_0=0.0,
+            phi_0=0.0,
+            omega_0=0.0,
+            lambda_0=0.0,
+            A_1=0.0,
+            phi_1=0.0,
+            omega_1=0.0,
+            lambda_1=0.0,
+        ),
+        np.array([1, 1, 1, 1, 1, 1], dtype=np.float64),
+    )
 
 
-def test_asymmetry_reducer_readable_signals_multiple_oscillator(
-    simpledae_multi_oscillator, multi_damped_oscillator_asymmetry_reducer
-):
-    assert set(
-        multi_damped_oscillator_asymmetry_reducer.additional_readable_signals(
-            simpledae_multi_oscillator
-        )
-    ) == {
-        simpledae_multi_oscillator.reducer.B,
-        simpledae_multi_oscillator.reducer.B_err,
-        simpledae_multi_oscillator.reducer.A_0,
-        simpledae_multi_oscillator.reducer.A_0_err,
-        simpledae_multi_oscillator.reducer.omega_0,
-        simpledae_multi_oscillator.reducer.omega_0_err,
-        simpledae_multi_oscillator.reducer.phi_0,
-        simpledae_multi_oscillator.reducer.phi_0_err,
-        simpledae_multi_oscillator.reducer.lambda_0,
-        simpledae_multi_oscillator.reducer.lambda_0_err,
-        simpledae_multi_oscillator.reducer.A_1,
-        simpledae_multi_oscillator.reducer.A_1_err,
-        simpledae_multi_oscillator.reducer.omega_1,
-        simpledae_multi_oscillator.reducer.omega_1_err,
-        simpledae_multi_oscillator.reducer.phi_1,
-        simpledae_multi_oscillator.reducer.phi_1_err,
-        simpledae_multi_oscillator.reducer.lambda_1,
-        simpledae_multi_oscillator.reducer.lambda_1_err,
+def test_asymmetry_reducer_readable_signals(simpledae, rebinning_asymmetry_reducer):
+    assert set(rebinning_asymmetry_reducer.additional_readable_signals(simpledae)) == {
+        simpledae.reducer.B,
+        simpledae.reducer.B_err,
+        simpledae.reducer.A_0,
+        simpledae.reducer.A_0_err,
+        simpledae.reducer.omega_0,
+        simpledae.reducer.omega_0_err,
+        simpledae.reducer.phi_0,
+        simpledae.reducer.phi_0_err,
+        simpledae.reducer.lambda_0,
+        simpledae.reducer.lambda_0_err,
     }
 
 
@@ -193,11 +157,9 @@ def test_rebin_and_sum_with_no_rebinning(asymmetry_reducer):
     scipp.testing.assert_allclose(result.coords["tof"], time)
 
 
-async def test_asymmetry_reducer(simpledae_single_oscillator):
-    simpledae_single_oscillator.trigger_and_get_specdata = AsyncMock(return_value=None)
-    simpledae_single_oscillator.reducer._first_det.read_spectrum_dataarray = AsyncMock(
-        return_value=None
-    )
+async def test_asymmetry_reducer(simpledae):
+    simpledae.trigger_and_get_specdata = AsyncMock(return_value=None)
+    simpledae.reducer._first_det.read_spectrum_dataarray = AsyncMock(return_value=None)
 
     with patch(
         "ibex_bluesky_core.devices.muon.MuonAsymmetryReducer._calculate_asymmetry",
@@ -215,10 +177,10 @@ async def test_asymmetry_reducer(simpledae_single_oscillator):
             },
         )
 
-        await simpledae_single_oscillator.reducer.reduce_data(simpledae_single_oscillator)
+        await simpledae.reducer.reduce_data(simpledae)
 
-    assert await simpledae_single_oscillator.reducer.B.get_value() == pytest.approx(0.0, abs=1e-8)
-    assert await simpledae_single_oscillator.reducer.A_0.get_value() == pytest.approx(0.0, abs=1e-8)
+    assert await simpledae.reducer.B.get_value() == pytest.approx(0.0, abs=1e-8)
+    assert await simpledae.reducer.A_0.get_value() == pytest.approx(0.0, abs=1e-8)
 
 
 async def test_asymmetry_reducer_real_data():
@@ -281,3 +243,54 @@ async def test_asymmetry_reducer_real_data():
     assert await dae.reducer.omega_0.get_value() == pytest.approx(omega_0, abs=1e-3)
     assert await dae.reducer.phi_0.get_value() == pytest.approx(phi_0, abs=1e-3)
     assert await dae.reducer.lambda_0.get_value() == pytest.approx(lambda_0, abs=1e-3)
+
+
+def test_missing_parameters():
+    with pytest.raises(ValueError, match=r"Missing parameters: .*"):
+        MuonAsymmetryReducer(
+            forward_detectors=np.array([1]),
+            backward_detectors=np.array([2]),
+            time_bin_edges=sc.linspace("tof", 0, 5, num=20, unit=sc.units.ns, dtype="float64"),
+            prefix="UNITTEST:",
+            model=damped_oscillator_model,
+            fit_parameters=lmfit.Parameters(),
+        )
+
+
+def test_calculate_asymmetry(asymmetry_reducer):
+    with (
+        patch.object(asymmetry_reducer, "_rebin_and_sum") as rebin_sum_mock,
+        patch("ibex_bluesky_core.devices.muon.calculate_polarisation") as calculate_asymmetry_mock,
+    ):
+        fake_da = MagicMock(spec=sc.DataArray)
+        rebin_sum_mock.return_value = sc.scalar(value=5.0, variance=10.0)
+        asymmetry_reducer._calculate_asymmetry(
+            np.array([[0], [0], [0]]), first_spec_dataarray=fake_da
+        )
+
+        calculate_asymmetry_mock.assert_called_with(
+            sc.scalar(value=5.0, variance=10.5),
+            sc.scalar(value=5.0, variance=10.5),
+            asymmetry_reducer._alpha,
+        )
+
+
+async def test_if_fit_fails_then_get_error(asymmetry_reducer, simpledae):
+    await asymmetry_reducer.connect(mock=True)
+
+    with (
+        patch.object(asymmetry_reducer, "_calculate_asymmetry"),
+        patch.object(asymmetry_reducer, "_fit_data") as fit_data_mock,
+        patch.object(simpledae, "trigger_and_get_specdata"),
+        patch.object(asymmetry_reducer._first_det, "read_spectrum_dataarray"),
+    ):
+        fit_data_mock.return_value = None
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "MuonAsymmetryReducer failed to fit asymmetry model to muon data.\n"
+                "Check beamline setup."
+            ),
+        ):
+            await asymmetry_reducer.reduce_data(simpledae)