Merge pull request #1130 from rhayes777/feature/graphical_analysis_output

feature/graphical analysis output

Author: Jammy2211

autofit/config/output.yaml

@@ -97,4 +97,4 @@ covariance: true # `covariance.csv`: The [free parameters x free parameters] cov
 data: true # `data.json`: The value of every data point in the data.
 noise_map: true # `noise_map.json`: The value of every RMS noise map value.
 
-search_log: true # `search.log`: logging produced whilst running the fit or fit_sequential method
+search_log: true # `search.log`: logging produced whilst running the fit method

autofit/graphical/declarative/abstract.py

@@ -7,7 +7,6 @@
 from autofit.graphical.declarative.graph import DeclarativeFactorGraph
 from autofit.graphical.expectation_propagation import AbstractFactorOptimiser
 from autofit.graphical.expectation_propagation import EPMeanField, EPOptimiser
-from autofit.mapper.model import ModelInstance
 from autofit.mapper.prior.abstract import Prior
 from autofit.mapper.prior_model.collection import Collection
 from autofit.mapper.variable import Plate
@@ -168,7 +167,7 @@ def optimise(
         optimiser: AbstractFactorOptimiser,
         paths: Optional[AbstractPaths] = None,
         ep_history: Optional = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Use an EP Optimiser to optimise the graph associated with this collection
@@ -198,29 +197,6 @@ def optimise(
             updated_ep_mean_field=updated_ep_mean_field,
         )
 
-    # TODO : Visualize method before fit?
-
-    def visualize(
-        self, paths: AbstractPaths, instance: ModelInstance, during_analysis: bool
-    ):
-        """
-        Visualise the instances provided using each factor.
-
-        Instances in the ModelInstance must have the same order as the factors.
-
-        Parameters
-        ----------
-        paths
-            Object describing where data should be saved to
-        instance
-            A collection of instances, each corresponding to a factor
-        during_analysis
-            Is this visualisation during analysis?
-        """
-        for model_factor, instance in zip(self.model_factors, instance):
-            model_factor.visualize(paths, instance, during_analysis)
-            model_factor.visualize_combined(None, paths, instance, during_analysis)
-
     @property
     def global_prior_model(self) -> Collection:
         """

autofit/graphical/declarative/collection.py

@@ -1,14 +1,17 @@
 from typing import Union, Optional
 
 from autofit.graphical.declarative.factor.hierarchical import HierarchicalFactor
-from autofit.mapper.model import ModelInstance
+
 from autofit.tools.namer import namer
 from .abstract import AbstractDeclarativeFactor
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.samples.pdf import SamplesPDF
 from autofit.non_linear.samples.summary import SamplesSummary
 from autofit.non_linear.analysis.combined import CombinedResult
 
+from autofit.mapper.model import ModelInstance
+from autofit.mapper.prior_model.prior_model import Model
+
 from autofit.jax_wrapper import register_pytree_node_class
 
 
@@ -165,3 +168,114 @@ def make_result(
             search_internal=search_internal,
             analysis=analysis,
         )
+
+    def _for_each_analysis(
+        self,
+        name,
+        paths,
+        *args,
+        **kwargs,
+    ):
+        """
+        Convenience function to call an underlying function for each
+        analysis with a paths object with an integer attached to the
+        end.
+
+        Parameters
+        ----------
+        func
+            Some function of the analysis class
+        paths
+            An object describing the paths for saving data (e.g. hard-disk directories or entries in sqlite database).
+        """
+        results = []
+        for (i, analysis), *args in zip(
+            enumerate(self.model_factors),
+            *args,
+        ):
+            child_paths = paths.for_sub_analysis(analysis_name=f"analyses/analysis_{i}")
+            func = getattr(analysis, name)
+            results.append(
+                func(
+                    child_paths,
+                    *args,
+                    **kwargs,
+                )
+            )
+
+        return results
+
+    def visualize(
+        self,
+        paths: AbstractPaths,
+        instance: ModelInstance,
+        during_analysis: bool,
+    ):
+        """
+        Visualise the instances provided using each factor.
+
+        Instances in the ModelInstance must have the same order as the factors.
+
+        Parameters
+        ----------
+        paths
+            Object describing where data should be saved to
+        instance
+            A collection of instances, each corresponding to a factor
+        during_analysis
+            Is this visualisation during analysis?
+        """
+        self._for_each_analysis(
+            "visualize",
+            paths,
+            instance,
+            during_analysis=during_analysis,
+        )
+
+    def visualize_before_fit(
+        self,
+        paths: AbstractPaths,
+        model: Model,
+    ):
+        """
+        Visualise the model provided using each factor.
+
+        Models in the ModelInstance must have the same order as the factors.
+
+        Parameters
+        ----------
+        paths
+            Object describing where data should be saved to
+        model
+            A collection of models, each corresponding to a factor
+        """
+        self._for_each_analysis(
+            "visualize_before_fit",
+            paths,
+            model,
+        )
+
+    def save_attributes(self, paths: AbstractPaths):
+        """
+        Save the attributes of the analysis to the paths object.
+        """
+        self._for_each_analysis("save_attributes", paths)
+
+    def save_results(self, paths: AbstractPaths, result):
+        """
+        Save the results of the analysis to the paths object.
+        """
+        self._for_each_analysis("save_results", paths, result)
+
+    def visualize_combined(
+        self,
+        instance,
+        paths: AbstractPaths,
+        during_analysis,
+    ):
+        self.model_factors[0].visualize_combined(
+            self.model_factors,
+            paths,
+            instance,
+            during_analysis=during_analysis,
+        )

autofit/graphical/declarative/factor/analysis.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+from autofit.mapper.prior_model.prior_model import Model
 from autofit.graphical.expectation_propagation import AbstractFactorOptimiser
 from autofit.mapper.model import ModelInstance
 from autofit.mapper.prior_model.prior_model import AbstractPriorModel
@@ -144,12 +145,49 @@ def visualize(
         self.analysis.visualize(
             paths=paths, instance=instance, during_analysis=during_analysis
         )
-        self.analysis.visualize_combined(
-            analyses=None,
-            paths=paths,
-            instance=instance,
-            during_analysis=during_analysis,
-        )
+
+    def visualize_before_fit(
+        self,
+        paths: AbstractPaths,
+        model: Model,
+    ):
+        """
+        Visualise the model provided using each factor.
+
+        Models in the ModelInstance must have the same order as the factors.
+
+        Parameters
+        ----------
+        paths
+            Object describing where data should be saved to
+        model
+            A collection of models, each corresponding to a factor
+        """
+        self.analysis.visualize_before_fit(paths=paths, model=model)
+
+    def save_attributes(self, paths: AbstractPaths):
+        """
+        Save the attributes of the analysis object to a file.
+
+        Parameters
+        ----------
+        paths
+            Object describing where data should be saved to
+        """
+        self.analysis.save_attributes(paths=paths)
+
+    def save_results(self, paths: AbstractPaths, result):
+        """
+        Save the results of the analysis to a file.
+
+        Parameters
+        ----------
+        paths
+            Object describing where data should be saved to
+        result
+            The result of the analysis
+        """
+        self.analysis.save_results(paths=paths, result=result)
 
     def log_likelihood_function(self, instance: ModelInstance) -> float:
         return self.analysis.log_likelihood_function(instance)