diff --git a/autofit/graphical/declarative/collection.py b/autofit/graphical/declarative/collection.py index f8ba1f291..c194c9f2d 100644 --- a/autofit/graphical/declarative/collection.py +++ b/autofit/graphical/declarative/collection.py @@ -302,28 +302,50 @@ def save_results(self, paths: AbstractPaths, result): """ self._for_each_analysis("save_results", paths, result) + def _factors_grouped_by_visualizer(self, instance): + """ + ``(factors, instances)`` pairs grouped by each factor's analysis ``Visualizer`` + class, order preserved within groups. + + A combined visualization can only draw analyses its ``Visualizer`` understands, + so a mixed-dataset graph (e.g. imaging + weak-lensing factors) must dispatch one + combined call per visualizer type rather than routing every factor into the lead + factor's visualizer. A homogeneous graph produces exactly one group, making the + dispatch identical to the previous lead-factor-takes-all behaviour. + """ + groups = {} + for factor, single_instance in zip(self.model_factors, instance): + analysis = getattr(factor, "analysis", factor) + key = getattr(type(analysis), "Visualizer", type(analysis)) + groups.setdefault(key, ([], [])) + groups[key][0].append(factor) + groups[key][1].append(single_instance) + return groups.values() + def visualize_combined( self, instance, paths: AbstractPaths, during_analysis, ): - self.model_factors[0].visualize_combined( - self.model_factors, - paths, - instance, - during_analysis=during_analysis, - ) + for factors, instances in self._factors_grouped_by_visualizer(instance): + factors[0].visualize_combined( + factors, + paths, + instances, + during_analysis=during_analysis, + ) def perform_quick_update(self, paths, instance): try: - self.model_factors[0].visualize_combined( - analyses=self.model_factors, - paths=paths, - instance=instance, - during_analysis=True, - quick_update=True, - ) + for factors, instances in self._factors_grouped_by_visualizer(instance): + factors[0].visualize_combined( + analyses=factors, + paths=paths, + instance=instances, + during_analysis=True, + quick_update=True, + ) except Exception as e: pass \ No newline at end of file diff --git a/test_autofit/graphical/test_visualize_combined_dispatch.py b/test_autofit/graphical/test_visualize_combined_dispatch.py new file mode 100644 index 000000000..f9e4bbed8 --- /dev/null +++ b/test_autofit/graphical/test_visualize_combined_dispatch.py @@ -0,0 +1,67 @@ +import autofit as af + + +class _RecordingVisualizerA(af.Visualizer): + calls = [] + + @staticmethod + def visualize_combined(analyses, paths, instance, during_analysis, quick_update=False): + _RecordingVisualizerA.calls.append((list(analyses), list(instance))) + + +class _RecordingVisualizerB(af.Visualizer): + calls = [] + + @staticmethod + def visualize_combined(analyses, paths, instance, during_analysis, quick_update=False): + _RecordingVisualizerB.calls.append((list(analyses), list(instance))) + + +class _AnalysisA(af.mock.MockAnalysis): + Visualizer = _RecordingVisualizerA + + +class _AnalysisB(af.mock.MockAnalysis): + Visualizer = _RecordingVisualizerB + + +def _factor_graph(analyses): + model = af.Model(af.ex.Gaussian) + factors = [af.AnalysisFactor(prior_model=model, analysis=analysis) for analysis in analyses] + return af.FactorGraphModel(*factors) + + +def test__homogeneous_graph__single_combined_call_with_all_factors(): + _RecordingVisualizerA.calls = [] + + analyses = [_AnalysisA(), _AnalysisA(), _AnalysisA()] + graph = _factor_graph(analyses) + + graph.visualize_combined(instance=["i0", "i1", "i2"], paths=None, during_analysis=True) + + assert len(_RecordingVisualizerA.calls) == 1 + called_analyses, called_instances = _RecordingVisualizerA.calls[0] + assert called_analyses == analyses + assert called_instances == ["i0", "i1", "i2"] + + +def test__mixed_graph__one_combined_call_per_visualizer_type(): + _RecordingVisualizerA.calls = [] + _RecordingVisualizerB.calls = [] + + analysis_a0, analysis_b, analysis_a1 = _AnalysisA(), _AnalysisB(), _AnalysisA() + graph = _factor_graph([analysis_a0, analysis_b, analysis_a1]) + + graph.visualize_combined(instance=["a0", "b0", "a1"], paths=None, during_analysis=True) + + # A-group: both A analyses, order preserved, with their matching instances. + assert len(_RecordingVisualizerA.calls) == 1 + called_analyses, called_instances = _RecordingVisualizerA.calls[0] + assert called_analyses == [analysis_a0, analysis_a1] + assert called_instances == ["a0", "a1"] + + # B-group: the lone B analysis with its own instance. + assert len(_RecordingVisualizerB.calls) == 1 + called_analyses, called_instances = _RecordingVisualizerB.calls[0] + assert called_analyses == [analysis_b] + assert called_instances == ["b0"]