From bbf8a280b5d6367362e4d0aff8cde5ecd23715e4 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 5 Nov 2026 14:51:37 +0100 Subject: [PATCH] fix: compute extension children on-the-fly to catch optional dependencies --- src/spikeinterface/core/sortinganalyzer.py | 66 ++++++++++--------- .../tests/test_analyzer_extension_core.py | 23 +++++-- 2 files changed, 53 insertions(+), 36 deletions(-) diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 1873b3134b..2ba1841ed2 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -2399,7 +2399,7 @@ def compute_one_extension(self, extension_name, save=True, verbose=False, **kwar save = False extension_class = get_extension_class(extension_name) - for child in _get_children_dependencies(extension_name): + for child in _get_children_dependencies(extension_name, self._get_extensions_params()): if self.has_extension(child): if verbose: print(f"Deleting extension: {child}") @@ -2497,8 +2497,9 @@ def compute_several_extensions( # Sort extensions by dependency order sorted_extensions = _sort_extensions_by_dependency(extensions) + all_extensions_params = {**self._get_extensions_params(), **extensions} for extension_name in sorted_extensions.keys(): - for child in _get_children_dependencies(extension_name): + for child in _get_children_dependencies(extension_name, all_extensions_params): if verbose: print(f"Deleting extension: {child}") self.delete_extension(child) @@ -2764,6 +2765,17 @@ def get_loaded_extension_names(self): """ return list(self.extensions.keys()) + def _get_extensions_params(self) -> dict[str, dict]: + """ + Return the params of all computed extensions (loaded or saved), keyed by extension name. + """ + # saved-but-unloaded extensions fall back to default params rather than being read from disk + params: dict[str, dict] = {} + if self.format != "memory": + params = {name: {} for name in self.get_saved_extension_names()} + params.update({name: ext.params for name, ext in self.extensions.items()}) + return params + def has_extension(self, extension_name: str) -> bool: """ Check if the extension exists in memory (dict) or in the folder or in zarr. @@ -2863,16 +2875,6 @@ def _sort_extensions_by_dependency(extensions): required_dependencies = get_extension_class(extension).get_required_dependencies(**extension_params) optional_dependencies = get_extension_class(extension).get_optional_dependencies(**extension_params) - for dependency in required_dependencies: - if "|" in dependency: - for dep in dependency.split("|"): - if extension not in _extension_children[dep]: - print(f"Adding missing dependency child link: {dep} -> {extension}") - _extension_children[dep].append(extension) - else: - if extension not in _extension_children[dependency]: - print(f"Adding missing dependency child link: {dependency} -> {extension}") - _extension_children[dependency].append(extension) # Split cases with an "or" in them, and flatten into a list all_dependencies = required_dependencies + optional_dependencies @@ -2905,11 +2907,8 @@ def _sort_extensions_by_dependency(extensions): global _possible_extensions _possible_extensions = [] -global _extension_children -_extension_children = {} - -def _get_children_dependencies(extension_name): +def _get_children_dependencies(extension_name, extensions): """ Extension classes have a `depend_on` attribute to declare on which class they depend on. For instance "templates" depends on "waveforms". "waveforms" depends on "random_spikes". @@ -2923,15 +2922,29 @@ def _get_children_dependencies(extension_name): For instance, recomputing the "waveforms" needs to delete the "templates", since the latter depends on the former. For this particular example, if we change the "ms_before" parameter of the "waveforms", also the "templates" will require recomputation as this parameter is inherited. + + Dependencies are resolved with the params of each extension, so optional dependencies (e.g. "quality_metrics" + on "spike_amplitudes") are included. Every alternative of a "a|b" dependency counts as a parent, since which + one was used is not recorded. + + Parameters + ---------- + extension_name : str + The parent extension name. + extensions : dict + Mapping of computed extension name -> params. + + Returns + ------- + names : list + The children, grand children, etc. of ``extension_name``. """ names = [] - children = _extension_children[extension_name] - for child in children: - if child not in names: + for child, params in extensions.items(): + if child not in names and extension_name in get_extension_class(child).get_any_dependencies(**params): names.append(child) - grand_children = _get_children_dependencies(child) - names.extend(grand_children) - return list(names) + names.extend(n for n in _get_children_dependencies(child, extensions) if n not in names) + return names def register_result_extension(extension_class): @@ -2958,15 +2971,6 @@ def register_result_extension(extension_class): _possible_extensions.append(extension_class) - # create the children dpendencies to be able to delete on re-compute - _extension_children[extension_class.extension_name] = [] - for parent_name in extension_class.get_required_dependencies(): - if "|" in parent_name: - for name in parent_name.split("|"): - _extension_children[name].append(extension_class.extension_name) - else: - _extension_children[parent_name].append(extension_class.extension_name) - def get_extension_class(extension_name: str, auto_import=True): """ diff --git a/src/spikeinterface/core/tests/test_analyzer_extension_core.py b/src/spikeinterface/core/tests/test_analyzer_extension_core.py index 6423ce11ce..9f6db410dd 100644 --- a/src/spikeinterface/core/tests/test_analyzer_extension_core.py +++ b/src/spikeinterface/core/tests/test_analyzer_extension_core.py @@ -8,7 +8,7 @@ from spikeinterface.core import create_sorting_analyzer from spikeinterface.core import Templates -from spikeinterface.core.sortinganalyzer import _extension_children, _get_children_dependencies +from spikeinterface.core.sortinganalyzer import _get_children_dependencies import numpy as np @@ -275,14 +275,27 @@ def test_ComputeNoiseLevels(format, sparse, create_cache_folder): def test_get_children_dependencies(): - assert "waveforms" in _extension_children["random_spikes"] - - rs_children = _get_children_dependencies("random_spikes") + extensions = {"random_spikes": {}, "waveforms": {}, "templates": {}} + rs_children = _get_children_dependencies("random_spikes", extensions) assert "waveforms" in rs_children assert "templates" in rs_children - assert rs_children.index("waveforms") < rs_children.index("templates") + # "templates" depends on "random_spikes|waveforms": both alternatives are parents + assert _get_children_dependencies("waveforms", extensions) == ["templates"] + # children not computed are not returned + assert _get_children_dependencies("random_spikes", {"random_spikes": {}}) == [] + + +def test_get_children_dependencies_optional(): + import spikeinterface.metrics # noqa: F401 registers quality_metrics + + # quality_metrics only optionally depends on spike_amplitudes, through its metrics + extensions = {"spike_amplitudes": {}, "quality_metrics": {"metric_names": ["amplitude_cutoff"]}} + assert _get_children_dependencies("spike_amplitudes", extensions) == ["quality_metrics"] + extensions["quality_metrics"] = {"metric_names": ["firing_rate"]} + assert _get_children_dependencies("spike_amplitudes", extensions) == [] + def test_delete_on_recompute(create_cache_folder): cache_folder = create_cache_folder