Fix and re-enable IRF Dispersion Test

Improve the tests in test_spectral_irf
Added NoDispersion test
Re-enable IrfDispersion tests on the CI (de-Volkswagen the tests)
Improve assignment of center_dispersion to result dataset

.github/workflows/CI_CD_actions.yml

@@ -109,7 +109,7 @@ jobs:
         run: pip freeze
       - name: Run tests
         run: |
-          pytest --cov=./ --cov-report term --cov-report xml --cov-config pyproject.toml -k 'not IrfDispersion' glotaran
+          pytest --cov=./ --cov-report term --cov-report xml --cov-config pyproject.toml glotaran
 
       - name: Codecov Upload
         uses: codecov/codecov-action@v2

glotaran/builtin/megacomplexes/decay/test/test_spectral_irf.py

@@ -1,4 +1,6 @@
+import warnings
 from copy import deepcopy
+from textwrap import dedent
 
 import numpy as np
 import pytest
@@ -35,6 +37,14 @@
     sh1:
         type: one
 """
+MODEL_NO_IRF_DISPERSION = f"""\
+{MODEL_BASE}
+irf:
+    irf1:
+        type: spectral-gaussian
+        center: irf.center
+        width: irf.width
+"""
 MODEL_SIMPLE_IRF_DISPERSION = f"""\
 {MODEL_BASE}
 irf:
@@ -57,13 +67,31 @@
         width_dispersion: [irf.width_dispersion]
 """
 
+MODEL_MULTIPULSE_IRF_DISPERSION = f"""\
+{MODEL_BASE}
+irf:
+    irf1:
+        type: spectral-multi-gaussian
+        center: [irf.center1, irf.center2]
+        width: [irf.width]
+        dispersion_center: irf.dispersion_center
+        center_dispersion: [irf.center_dispersion1, irf.center_dispersion2, irf.center_dispersion3]
+"""
+
 PARAMETERS_BASE = """\
 j:
     - ['1', 1, {'vary': False, 'non-negative': False}]
 kinetic:
     - ['1', 0.5, {'non-negative': False}]
 """
 
+PARAMETERS_NO_IRF_DISPERSION = f"""\
+{PARAMETERS_BASE}
+irf:
+    - ['center', 0.3]
+    - ['width', 0.1]
+"""
+
 PARAMETERS_SIMPLE_IRF_DISPERSION = f"""\
 {PARAMETERS_BASE}
 irf:
@@ -73,52 +101,92 @@
     - ['center_dispersion', 0.5]
 """
 
+# What is this?
 PARAMETERS_MULTI_IRF_DISPERSION = f"""\
 {PARAMETERS_BASE}
 irf:
     - ["center", 0.3]
     - ["width", 0.1]
     - ["dispersion_center", 400, {{"vary": False}}]
-    - ["center_dispersion1", 0.01]
-    - ["center_dispersion2", 0.001]
+    - ["center_dispersion1", 0.1]
+    - ["center_dispersion2", 0.01]
     - ["width_dispersion", 0.025]
 """
 
+PARAMETERS_MULTIPULSE_IRF_DISPERSION = f"""\
+{PARAMETERS_BASE}
+irf:
+    - ["center1", 0.3]
+    - ["center2", 0.4]
+    - ['width', 0.1]
+    - ['dispersion_center', 400, {{'vary': False}}]
+    - ["center_dispersion1", 0.5]
+    - ["center_dispersion2", 0.1]
+    - ["center_dispersion3", -0.01]
+"""
+
+
+def _time_axis():
+    time_p1 = np.linspace(-1, 1, 20, endpoint=False)
+    time_p2 = np.linspace(1, 2, 10, endpoint=False)
+    time_p3 = np.geomspace(2, 20, num=20)
+    return np.array(np.concatenate([time_p1, time_p2, time_p3]))
+
+
+def _spectral_axis():
+    return np.linspace(300, 500, 3)
+
+
+def _calculate_irf_position(index, center, dispersion_center=None, center_dispersion=None):
+    if center_dispersion is None:
+        center_dispersion = []
+    if dispersion_center is not None:
+        distance = (index - dispersion_center) / 100
+    if dispersion_center is not None:
+        for i, coefficient in enumerate(center_dispersion):
+            center += coefficient * np.power(distance, i + 1)
+    return center
+
+
+class NoIrfDispersion:
+    model = load_model(MODEL_NO_IRF_DISPERSION, format_name="yml_str")
+    parameters = load_parameters(PARAMETERS_NO_IRF_DISPERSION, format_name="yml_str")
+    axis = {"time": _time_axis(), "spectral": _spectral_axis()}
+
 
 class SimpleIrfDispersion:
     model = load_model(MODEL_SIMPLE_IRF_DISPERSION, format_name="yml_str")
     parameters = load_parameters(PARAMETERS_SIMPLE_IRF_DISPERSION, format_name="yml_str")
-    time_p1 = np.linspace(-1, 2, 50, endpoint=False)
-    time_p2 = np.linspace(2, 5, 30, endpoint=False)
-    time_p3 = np.geomspace(5, 10, num=20)
-    time = np.array(np.concatenate([time_p1, time_p2, time_p3]))
-    spectral = np.arange(300, 500, 100)
-    axis = {"time": time, "spectral": spectral}
+    axis = {"time": _time_axis(), "spectral": _spectral_axis()}
 
 
 class MultiIrfDispersion:
     model = load_model(MODEL_MULTI_IRF_DISPERSION, format_name="yml_str")
     parameters = load_parameters(PARAMETERS_MULTI_IRF_DISPERSION, format_name="yml_str")
-    time = np.arange(-1, 5, 0.2)
-    spectral = np.arange(300, 500, 100)
-    axis = {"time": time, "spectral": spectral}
+    axis = {"time": _time_axis(), "spectral": _spectral_axis()}
+
+
+class MultiCenterIrfDispersion:
+    model = load_model(MODEL_MULTIPULSE_IRF_DISPERSION, format_name="yml_str")
+    parameters = load_parameters(PARAMETERS_MULTIPULSE_IRF_DISPERSION, format_name="yml_str")
+    axis = {"time": _time_axis(), "spectral": _spectral_axis()}
 
 
 @pytest.mark.parametrize(
     "suite",
     [
+        NoIrfDispersion,
         SimpleIrfDispersion,
         MultiIrfDispersion,
+        MultiCenterIrfDispersion,
     ],
 )
 def test_spectral_irf(suite):
 
     model = suite.model
-    print(model.validate())
     assert model.valid()
 
     parameters = suite.parameters
-    print(model.validate(parameters))
     assert model.valid(parameters)
 
     sim_model = deepcopy(model)
@@ -136,27 +204,60 @@ def test_spectral_irf(suite):
         maximum_number_function_evaluations=20,
     )
     result = optimize(scheme)
-    print(result.optimized_parameters)
+    # print(result.optimized_parameters)
 
     for label, param in result.optimized_parameters.all():
         assert np.allclose(param.value, parameters.get(label).value, rtol=1e-1)
 
     resultdata = result.data["dataset1"]
 
-    print(resultdata)
-
+    # print(resultdata)
     assert np.array_equal(dataset["time"], resultdata["time"])
     assert np.array_equal(dataset["spectral"], resultdata["spectral"])
     assert dataset.data.shape == resultdata.data.shape
     assert dataset.data.shape == resultdata.fitted_data.shape
-    assert np.allclose(dataset.data, resultdata.fitted_data, atol=1e-14)
+    # assert np.allclose(dataset.data, resultdata.fitted_data, atol=1e-14)
+
+    fit_data_max_at_start = resultdata.fitted_data.isel(spectral=0).argmax(axis=0)
+    fit_data_max_at_end = resultdata.fitted_data.isel(spectral=-1).argmax(axis=0)
+
+    if suite is NoIrfDispersion:
+        assert "center_dispersion_1" not in resultdata
+        assert fit_data_max_at_start == fit_data_max_at_end
+    else:
+        assert "center_dispersion_1" in resultdata
+        assert fit_data_max_at_start != fit_data_max_at_end
+        if abs(fit_data_max_at_start - fit_data_max_at_end) < 3:
+            warnings.warn(
+                dedent(
+                    """
+                    Bad test, one of the following could be the case:
+                    - dispersion too small
+                    - spectral window to small
+                    - time resolution (around the maximum of the IRF) too low"
+                    """
+                )
+            )
 
-    irf_max_at_start = resultdata.fitted_data.isel(spectral=0).argmax(axis=0)
-    irf_max_at_end = resultdata.fitted_data.isel(spectral=-1).argmax(axis=0)
-    print(f" irf_max_at_start: {irf_max_at_start}\n irf_max_at_end: {irf_max_at_end}")
-    # These should not be equal due to dispersion:
-    assert irf_max_at_start != irf_max_at_end
+        for x in suite.axis["spectral"]:
+            # calculated irf location
+            model_irf_center = suite.model.irf["irf1"].center
+            model_dispersion_center = suite.model.irf["irf1"].dispersion_center
+            model_center_dispersion = suite.model.irf["irf1"].center_dispersion
+            calc_irf_location_at_x = _calculate_irf_position(
+                x, model_irf_center, model_dispersion_center, model_center_dispersion
+            )
+            # fitted irf location
+            fitted_irf_loc_at_x = resultdata["center_dispersion"].sel(spectral=x)
+            assert np.allclose(calc_irf_location_at_x, fitted_irf_loc_at_x)
 
     assert "species_associated_spectra" in resultdata
     assert "decay_associated_spectra" in resultdata
     assert "irf_center" in resultdata
+
+
+if __name__ == "__main__":
+    test_spectral_irf(NoIrfDispersion)
+    test_spectral_irf(SimpleIrfDispersion)
+    test_spectral_irf(MultiIrfDispersion)
+    test_spectral_irf(MultiCenterIrfDispersion)

glotaran/builtin/megacomplexes/decay/util.py

@@ -228,10 +228,10 @@ def retrieve_irf(dataset_model: DatasetModel, dataset: xr.Dataset, global_dimens
     dataset["irf_center"] = ("irf_nr", center) if len(center) > 1 else center[0]
     dataset["irf_width"] = ("irf_nr", width) if len(width) > 1 else width[0]
     if isinstance(irf, IrfSpectralMultiGaussian) and irf.dispersion_center:
-        for i, dispersion in enumerate(
-            irf.calculate_dispersion(dataset.coords["spectral"].values)
-        ):
-            dataset[f"center_dispersion_{i+1}"] = (
-                global_dimension,
-                dispersion,
-            )
+        # TODO: rename center_dispersion to irf_center_location
+        dataset["center_dispersion"] = (
+            ("irf_nr", global_dimension),
+            irf.calculate_dispersion(dataset.coords["spectral"].values),
+        )
+        # TODO: deprecate and remove old _n naming for 'center_dispersion'
+        dataset["center_dispersion_1"] = dataset["center_dispersion"].sel(irf_nr=0)