Move index dependent calculation to megacomplexes for speed-up

.github/workflows/CI_CD_actions.yml

@@ -55,7 +55,7 @@ jobs:
         run: |
           conda install -y pandoc
           python -m pip install -U pip wheel
-          pip install .
+          python -m pip install .
           python -m pip install -U -r docs/requirements.txt
       - name: Show installed packages
         run: pip freeze
@@ -80,7 +80,7 @@ jobs:
         run: |
           conda install -y pandoc
           python -m pip install -U pip wheel
-          pip install .
+          python -m pip install .
           python -m pip install -U -r docs/requirements.txt
       - name: Show installed packages
         run: pip freeze
@@ -106,10 +106,10 @@ jobs:
           python -m pip install -U -r requirements_dev.txt
           pip install .
       - name: Show installed packages
-        run: pip freeze
+        run: python -m pip freeze
       - name: Build docs
         run: |
-          py.test -vv --nbval docs/source/notebooks
+          python -m pytest -vv --nbval docs/source/notebooks
 
   test:
     runs-on: ${{ matrix.os }}
@@ -129,12 +129,12 @@ jobs:
         run: |
           python -m pip install --upgrade pip wheel
           python -m pip install -r requirements_dev.txt
-          pip install -e .
+          python -m pip install -e .
       - name: Show installed packages
-        run: pip freeze
+        run: python -m pip freeze
       - name: Run tests
         run: |
-          pytest --cov=./ --cov-report term --cov-report xml --cov-config pyproject.toml glotaran
+          python -m pytest --cov=./ --cov-report term --cov-report xml --cov-config pyproject.toml glotaran
 
       - name: Codecov Upload
         uses: codecov/codecov-action@v3

benchmark/pytest/analysis/test_optimization_group.py

@@ -44,7 +44,6 @@ class BenchmarkMegacomplex(Megacomplex):
     def calculate_matrix(
         self,
         dataset_model,
-        global_index: int | None,
         global_axis: np.typing.ArrayLike,
         model_axis: np.typing.ArrayLike,
         **kwargs,

changelog.md

@@ -12,6 +12,7 @@
 
 - ✨ Add optimization history to result and iteration column to parameter history (#1134)
 - ♻️ Complete refactor of model and parameter packages using attrs (#1135)
+- ♻️ Move index dependent calculation to megacomplexes for speed-up (#1175)
 
 ### 👌 Minor Improvements:
 

glotaran/builtin/megacomplexes/baseline/baseline_megacomplex.py

@@ -15,7 +15,6 @@ class BaselineMegacomplex(Megacomplex):
     def calculate_matrix(
         self,
         dataset_model: DatasetModel,
-        global_index: int | None,
         global_axis: np.typing.ArrayLike,
         model_axis: np.typing.ArrayLike,
         **kwargs,
@@ -24,9 +23,6 @@ def calculate_matrix(
         matrix = np.ones((model_axis.size, 1), dtype=np.float64)
         return clp_label, matrix
 
-    def index_dependent(self, dataset_model: DatasetModel) -> bool:
-        return False
-
     def finalize_data(
         self,
         dataset_model: DatasetModel,

glotaran/builtin/megacomplexes/baseline/test/test_baseline_megacomplex.py

@@ -19,7 +19,7 @@ def test_baseline():
     time = np.asarray(np.arange(0, 50, 1.5))
     pixel = np.asarray([0])
     dataset_model = fill_item(model.dataset["dataset1"], model, parameters)
-    matrix = MatrixProvider.calculate_dataset_matrix(dataset_model, None, pixel, time)
+    matrix = MatrixProvider.calculate_dataset_matrix(dataset_model, pixel, time)
     compartments = matrix.clp_labels
 
     assert len(compartments) == 1

glotaran/builtin/megacomplexes/clp_guide/clp_guide_megacomplex.py

@@ -16,7 +16,6 @@ class ClpGuideMegacomplex(Megacomplex):
     def calculate_matrix(
         self,
         dataset_model: DatasetModel,
-        global_index: int | None,
         global_axis: np.typing.ArrayLike,
         model_axis: np.typing.ArrayLike,
         **kwargs,
@@ -25,9 +24,6 @@ def calculate_matrix(
         matrix = np.ones((1, 1), dtype=np.float64)
         return clp_label, matrix
 
-    def index_dependent(self, dataset_model: DatasetModel) -> bool:
-        return False
-
     def finalize_data(
         self,
         dataset_model: DatasetModel,

glotaran/builtin/megacomplexes/coherent_artifact/coherent_artifact_megacomplex.py

@@ -13,7 +13,6 @@
 from glotaran.model import Megacomplex
 from glotaran.model import ModelError
 from glotaran.model import ParameterType
-from glotaran.model import is_dataset_model_index_dependent
 from glotaran.model import megacomplex
 
 
@@ -27,35 +26,60 @@ class CoherentArtifactMegacomplex(Megacomplex):
     def calculate_matrix(
         self,
         dataset_model: DatasetModel,
-        global_index: int | None,
         global_axis: np.typing.ArrayLike,
         model_axis: np.typing.ArrayLike,
         **kwargs,
     ):
         if not 1 <= self.order <= 3:
             raise ModelError("Coherent artifact order must be between in [1,3]")
 
-        if dataset_model.irf is None:
+        irf = dataset_model.irf
+        if irf is None:
             raise ModelError(f'No irf in dataset "{dataset_model.label}"')
 
-        if not isinstance(dataset_model.irf, IrfMultiGaussian):
+        if not isinstance(irf, IrfMultiGaussian):
             raise ModelError(f'Irf in dataset "{dataset_model.label} is not a gaussian irf."')
 
-        irf = dataset_model.irf
+        matrix_shape = (
+            (global_axis.size, model_axis.size, self.order)
+            if index_dependent(dataset_model)
+            else (model_axis.size, self.order)
+        )
+        matrix = np.zeros(matrix_shape, dtype=np.float64)
+        if index_dependent(dataset_model):
+            centers, widths = [], []
+            for global_index in range(global_axis.size):
+                center, width = self.get_irf_parameter(irf, global_index, global_axis)
+                centers.append(center)
+                widths.append(width)
+            _calculate_coherent_artifact_matrix(
+                matrix,
+                np.asarray(centers),
+                np.asarray(widths),
+                global_axis.size,
+                model_axis,
+                self.order,
+            )
+
+        else:
+            center, width = self.get_irf_parameter(irf, None, global_axis)
+            _calculate_coherent_artifact_matrix_on_index(
+                matrix, center, width, model_axis, self.order
+            )
+
+        return self.compartments(), matrix
 
+    def get_irf_parameter(
+        self, irf: IrfMultiGaussian, global_index: int | None, global_axis: np.typing.ArrayLike
+    ) -> tuple[float, float]:
         center, width, _, shift, _, _ = irf.parameter(global_index, global_axis)
         center = center[0] - shift
         width = self.width.value if self.width is not None else width[0]
-
-        matrix = _calculate_coherent_artifact_matrix(center, width, model_axis, self.order)
-        return self.compartments(), matrix
+        return center, width
 
     def compartments(self):
         return [f"coherent_artifact_{i}" for i in range(1, self.order + 1)]
 
-    def index_dependent(self, dataset_model: DatasetModel) -> bool:
-        return index_dependent(dataset_model)
-
     def finalize_data(
         self,
         dataset_model: DatasetModel,
@@ -68,7 +92,7 @@ def finalize_data(
             model_dimension = dataset.attrs["model_dimension"]
             dataset.coords["coherent_artifact_order"] = np.arange(1, self.order + 1)
             response_dimensions = (model_dimension, "coherent_artifact_order")
-            if is_dataset_model_index_dependent(dataset_model):
+            if len(dataset.matrix.shape) == 3:
                 response_dimensions = (global_dimension, *response_dimensions)
             dataset["coherent_artifact_response"] = (
                 response_dimensions,
@@ -81,9 +105,18 @@ def finalize_data(
         retrieve_irf(dataset_model, dataset, global_dimension)
 
 
+@nb.jit(nopython=True, parallel=False)
+def _calculate_coherent_artifact_matrix(
+    matrix, centers, widths, global_axis_size, model_axis, order
+):
+    for i in nb.prange(global_axis_size):
+        _calculate_coherent_artifact_matrix_on_index(
+            matrix[i], centers[i], widths[i], model_axis, order
+        )
+
+
 @nb.jit(nopython=True, parallel=True)
-def _calculate_coherent_artifact_matrix(center, width, axis, order):
-    matrix = np.zeros((axis.size, order), dtype=np.float64)
+def _calculate_coherent_artifact_matrix_on_index(matrix, center, width, axis, order):
 
     matrix[:, 0] = np.exp(-1 * (axis - center) ** 2 / (2 * width**2))
     if order > 1:
@@ -93,4 +126,3 @@ def _calculate_coherent_artifact_matrix(center, width, axis, order):
         matrix[:, 2] = (
             matrix[:, 0] * (center**2 - width**2 - 2 * center * axis + axis**2) / width**4
         )
-    return matrix

glotaran/builtin/megacomplexes/coherent_artifact/test/test_coherent_artifact.py

@@ -88,15 +88,18 @@ def test_coherent_artifact(spectral_dependence: str):
     spectral = np.asarray([200, 300, 400])
 
     dataset_model = fill_item(model.dataset["dataset1"], model, parameters)
-    matrix = MatrixProvider.calculate_dataset_matrix(dataset_model, 0, spectral, time)
+    matrix = MatrixProvider.calculate_dataset_matrix(dataset_model, spectral, time)
     compartments = matrix.clp_labels
 
     print(compartments)
     assert len(compartments) == 4
     for i in range(1, 4):
         assert compartments[i] == f"coherent_artifact_{i}"
 
-    assert matrix.matrix.shape == (time.size, 4)
+    if spectral_dependence == "none":
+        assert matrix.matrix.shape == (time.size, 4)
+    else:
+        assert matrix.matrix.shape == (spectral.size, time.size, 4)
 
     clp = xr.DataArray(
         np.ones((3, 4)),
@@ -113,7 +116,7 @@ def test_coherent_artifact(spectral_dependence: str):
             ),
         ],
     )
-    axis = {"time": time, "spectral": clp.spectral}
+    axis = {"time": time, "spectral": clp.spectral.data}
     data = simulate(model, "dataset1", parameters, axis, clp)
 
     dataset = {"dataset1": data}

glotaran/builtin/megacomplexes/damped_oscillation/damped_oscillation_megacomplex.py

@@ -7,6 +7,7 @@
 
 from glotaran.builtin.megacomplexes.decay.decay_parallel_megacomplex import DecayDatasetModel
 from glotaran.builtin.megacomplexes.decay.irf import IrfMultiGaussian
+from glotaran.builtin.megacomplexes.decay.util import index_dependent
 from glotaran.model import DatasetModel
 from glotaran.model import ItemIssue
 from glotaran.model import Megacomplex
@@ -67,7 +68,6 @@ class DampedOscillationMegacomplex(Megacomplex):
     def calculate_matrix(
         self,
         dataset_model: DatasetModel,
-        global_index: int | None,
         global_axis: np.typing.ArrayLike,
         model_axis: np.typing.ArrayLike,
         **kwargs,
@@ -88,34 +88,29 @@ def calculate_matrix(
         )
         rates = np.array(self.rates)
 
-        matrix = np.ones((model_axis.size, len(clp_label)), dtype=np.float64)
+        irf = dataset_model.irf
+        matrix_shape = (
+            (global_axis.size, model_axis.size, len(clp_label))
+            if index_dependent(dataset_model)
+            else (model_axis.size, len(clp_label))
+        )
+        matrix = np.ones(matrix_shape, dtype=np.float64)
 
-        if dataset_model.irf is None:
+        if irf is None:
             calculate_damped_oscillation_matrix_no_irf(matrix, frequencies, rates, model_axis)
-        elif isinstance(dataset_model.irf, IrfMultiGaussian):
-            centers, widths, scales, shift, _, _ = dataset_model.irf.parameter(
-                global_index, global_axis
-            )
-            for center, width, scale in zip(centers, widths, scales):
-                matrix += calculate_damped_oscillation_matrix_gaussian_irf(
-                    frequencies,
-                    rates,
-                    model_axis,
-                    center,
-                    width,
-                    shift,
-                    scale,
+        elif isinstance(irf, IrfMultiGaussian):
+            if index_dependent(dataset_model):
+                for i in range(global_axis.size):
+                    calculate_damped_oscillation_matrix_gaussian_irf_on_index(
+                        matrix[i], frequencies, rates, irf, i, global_axis, model_axis
+                    )
+            else:
+                calculate_damped_oscillation_matrix_gaussian_irf_on_index(
+                    matrix, frequencies, rates, irf, None, global_axis, model_axis
                 )
-            matrix /= np.sum(scales)
 
         return clp_label, matrix
 
-    def index_dependent(self, dataset_model: DatasetModel) -> bool:
-        return (
-            isinstance(dataset_model.irf, IrfMultiGaussian)
-            and dataset_model.irf.is_index_dependent()
-        )
-
     def finalize_data(
         self,
         dataset_model: DatasetModel,
@@ -159,7 +154,7 @@ def finalize_data(
             phase,
         )
 
-        if self.index_dependent(dataset_model):
+        if index_dependent(dataset_model):
             dataset[f"{prefix}_sin"] = (
                 (
                     global_dimension,
@@ -200,6 +195,29 @@ def calculate_damped_oscillation_matrix_no_irf(matrix, frequencies, rates, axis)
         idx += 2
 
 
+def calculate_damped_oscillation_matrix_gaussian_irf_on_index(
+    matrix: np.typing.ArrayLike,
+    frequencies: np.typing.ArrayLike,
+    rates: np.typing.ArrayLike,
+    irf: IrfMultiGaussian,
+    global_index: int | None,
+    global_axis: np.typing.ArrayLike,
+    model_axis: np.typing.ArrayLike,
+):
+    centers, widths, scales, shift, _, _ = irf.parameter(global_index, global_axis)
+    for center, width, scale in zip(centers, widths, scales):
+        matrix += calculate_damped_oscillation_matrix_gaussian_irf(
+            frequencies,
+            rates,
+            model_axis,
+            center,
+            width,
+            shift,
+            scale,
+        )
+    matrix /= np.sum(scales)
+
+
 def calculate_damped_oscillation_matrix_gaussian_irf(
     frequencies: np.ndarray,
     rates: np.ndarray,